app.py

import gradio as gr
import spaces
import torch
from model import ModelHandler
from generator import Generator

# 1. Initialize Models Globally (in RAM)
# ZeroGPU will move them to VRAM inside the @spaces.GPU function
print("Initializing Application...")
handler = ModelHandler()
handler.load_models()
gen = Generator(handler)

# 2. Define GPU-enabled Inference Function
@spaces.GPU(duration=35)
def process_img(
    image, 
    prompt, 
    cfg_scale, 
    steps, 
    img_strength, 
    depth_strength, 
    edge_strength
):
    if image is None:
        raise gr.Error("Please upload an image first.")
    
    try:
        print("--- Starting Generation ---")
        # Pass all parameters to the generator
        result = gen.predict(
            image, 
            prompt,
            guidance_scale=cfg_scale,
            num_inference_steps=steps,
            img2img_strength=img_strength,
            depth_strength=depth_strength,
            lineart_strength=edge_strength
        )
        print("--- Generation Complete ---")
        return result
        
    except Exception as e:
        print(f"Error during generation: {e}")
        raise gr.Error(f"An error occurred: {str(e)}")

# 3. Build Gradio Interface
with gr.Blocks(title="Face To Pixel Art", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎮 Face to Pixel Art
        Upload any image. If there is a face, we'll keep the identity. If not, we'll pixelate the scene!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            input_img = gr.Image(type="pil", label="Input Image")
            prompt = gr.Textbox(
                label="Prompt (Optional)", 
                placeholder="Leave empty for auto-captioning...",
                info="The trigger words 'p1x3l4rt, pixel art' are added automatically."
            )
            
            # --- MOVED ACCORDION HERE ---
            with gr.Accordion("Advanced Settings", open=False):
                cfg_scale = gr.Slider(
                    elem_id="cfg_scale",
                    minimum=1.0, 
                    maximum=5.0, 
                    step=0.1, 
                    value=1.5, 
                    label="CFG Scale"
                )
                steps = gr.Slider(
                    elem_id="steps",
                    minimum=4, 
                    maximum=20, 
                    step=1, 
                    value=6, 
                    label="Steps Number"
                )
                img_strength = gr.Slider(
                    elem_id="img_strength",
                    minimum=0.1, 
                    maximum=1.0, 
                    step=0.05, 
                    value=0.3, 
                    label="Image Strength (Img2Img)"
                )
                depth_strength = gr.Slider(
                    elem_id="depth_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=0.3, 
                    label="DepthMap Strength"
                )
                edge_strength = gr.Slider(
                    elem_id="edge_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=0.3, 
                    label="EdgeMap Strength (LineArt)"
                )
            # --- END OF MOVED BLOCK ---

            run_btn = gr.Button("Generate Pixel Art", variant="primary")
            
            # --- FIX: Pass component OBJECTS, not strings ---
            gr.Examples(
                examples=[
                    # Add default values for new sliders
                    ["./examples/monalisa.jpg", "mona lisa portrait", 1.5, 6, 0.3, 0.3, 0.3],
                    ["./examples/man.jpg", "photo of a man", 1.5, 6, 0.3, 0.3, 0.3],
                    ["./examples/landscape.jpg", "a beautiful landscape", 1.5, 6, 0.3, 0.3, 0.3]
                ],
                # Pass the actual component variables
                inputs=[
                    input_img, 
                    prompt, 
                    cfg_scale, 
                    steps, 
                    img_strength, 
                    depth_strength, 
                    edge_strength
                ]
            )
            
        with gr.Column(scale=1):
            output_img = gr.Image(label="Pixel Art Result")
            
            # The accordion was here, but has been moved to the left column
                
    # Event Handler
    # Update inputs list for the click event
    all_inputs = [
        input_img, 
        prompt, 
        cfg_scale, 
        steps, 
        img_strength, 
        depth_strength, 
        edge_strength
    ]
    
    run_btn.click(
        fn=process_img, 
        inputs=all_inputs, 
        outputs=[output_img]
    )


# 4. Launch the App
if __name__ == "__main__":
    demo.queue(max_size=20, api_open=True)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_api=True
    )