app.py

import gradio as gr
import torch
from PIL import Image
import base64
from io import BytesIO
import json

print("🚀 Starting Affecto Inference Service...")

# Import our MagicFace model
from magicface_model import MagicFaceModel

# Initialize model
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🖥️  Device: {device}")

print("📥 Loading MagicFace model...")
model = MagicFaceModel(device=device)
print("✅ Model ready!")

# ============================================
# UTILITY FUNCTIONS
# ============================================

def pil_to_base64(image):
    """Convert PIL to base64"""
    buffered = BytesIO()
    image.save(buffered, format="JPEG", quality=95)
    return base64.b64encode(buffered.getvalue()).decode()

def base64_to_pil(base64_str):
    """Convert base64 to PIL"""
    image_bytes = base64.b64decode(base64_str)
    return Image.open(BytesIO(image_bytes))

# ============================================
# INFERENCE FUNCTIONS
# ============================================

def transform_gradio(image, au_params_str):
    """Gradio interface function"""
    try:
        # Parse AU params
        au_params = json.loads(au_params_str)
        
        # Ensure image is 512x512
        if image.size != (512, 512):
            image = image.resize((512, 512), Image.LANCZOS)
        
        # Transform
        result_image = model.transform(image, au_params)
        
        return result_image
    except Exception as e:
        print(f"❌ Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return image

def transform_api(image_base64, au_params_str):
    """API function for external calls"""
    try:
        print(f"📥 Received API request")
        
        # Decode image
        image = base64_to_pil(image_base64)
        print(f"📸 Image size: {image.size}")
        
        # Parse AU params
        au_params = json.loads(au_params_str)
        
        # Ensure 512x512
        if image.size != (512, 512):
            image = image.resize((512, 512), Image.LANCZOS)
        
        # Transform
        result_image = model.transform(image, au_params)
        
        # Encode result
        result_base64 = pil_to_base64(result_image)
        
        print("✅ Transformation complete")
        
        return result_base64
    except Exception as e:
        print(f"❌ API Error: {str(e)}")
        import traceback
        traceback.print_exc()
        raise

# ============================================
# GRADIO INTERFACE
# ============================================

with gr.Blocks(theme=gr.themes.Soft(), title="Affecto MagicFace API") as demo:
    gr.Markdown("# 🎭 Affecto - MagicFace Emotion Transformation")
    gr.Markdown("Transform facial emotions using Action Units (AU)")
    gr.Markdown("⚠️ **Note:** Currently using simplified model. Full MagicFace pipeline coming soon!")
    
    with gr.Tab("🖼️ Web Interface"):
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Upload Face Image (512x512 recommended)")
                au_params_input = gr.Textbox(
                    label="AU Parameters (JSON)",
                    value='{"AU6": 2.0, "AU12": 2.0}',
                    lines=3
                )
                transform_btn = gr.Button("✨ Transform", variant="primary", size="lg")
            
            with gr.Column():
                output_image = gr.Image(type="pil", label="Transformed Result")
        
        gr.Markdown("### 🎨 Emotion Presets (click to use):")
        gr.Examples(
            examples=[
                ['{"AU6": 2.0, "AU12": 2.0}'],  # Happy
                ['{"AU1": 2.0, "AU4": 2.0, "AU15": 2.0}'],  # Sad
                ['{"AU4": 3.0, "AU5": 2.0, "AU7": 2.0}'],  # Angry
                ['{"AU1": 3.0, "AU2": 2.0, "AU5": 3.0, "AU26": 2.0}'],  # Surprised
            ],
            inputs=[au_params_input],
            label="Emotion Presets"
        )
        
        transform_btn.click(
            fn=transform_gradio,
            inputs=[input_image, au_params_input],
            outputs=output_image
        )
    
    with gr.Tab("📡 API Documentation"):
        gr.Markdown("""
        ## API Usage
        
        ### Gradio API Endpoint
        
        ```python
        import requests
        import base64
        import json
        
        # Prepare image
        with open("face.jpg", "rb") as f:
            image_base64 = base64.b64encode(f.read()).decode()
        
        # Call API
        response = requests.post(
            "https://gauravvjhaa-affecto-inference.hf.space/api/predict",
            json={
                "data": [
                    image_base64,
                    '{"AU6": 2.0, "AU12": 2.0}'
                ]
            }
        )
        
        result = response.json()
        result_image = result["data"][0]  # base64 string
        ```
        
        ### Available Action Units:
        - **AU1** (0): Inner Brow Raiser - Values: 0-4
        - **AU2** (1): Outer Brow Raiser - Values: 0-4
        - **AU4** (2): Brow Lowerer - Values: 0-4
        - **AU5** (3): Upper Lid Raiser - Values: 0-4
        - **AU6** (4): Cheek Raiser - Values: 0-4
        - **AU9** (5): Nose Wrinkler - Values: 0-4
        - **AU12** (6): Lip Corner Puller (Smile) - Values: 0-4
        - **AU15** (7): Lip Corner Depressor - Values: 0-4
        - **AU17** (8): Chin Raiser - Values: 0-4
        - **AU20** (9): Lip Stretcher - Values: 0-4
        - **AU25** (10): Lips Part - Values: 0-4
        - **AU26** (11): Jaw Drop - Values: 0-4
        
        ### Example Combinations:
        - **Happy**: `{"AU6": 2, "AU12": 2}`
        - **Sad**: `{"AU1": 2, "AU4": 2, "AU15": 2}`
        - **Angry**: `{"AU4": 3, "AU5": 2, "AU7": 2}`
        - **Surprised**: `{"AU1": 3, "AU2": 2, "AU5": 3, "AU26": 2}`
        """)

print("✅ Affecto MagicFace API Ready!")
print(f"🌐 Gradio UI: https://gauravvjhaa-affecto-inference.hf.space/")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)