app.py

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

# 1. Initialize Models Globally
print("Initializing Application...")
handler = ModelHandler()
handler.load_models()
gen = Generator(handler)

# 2. Define GPU-enabled Inference Function
@spaces.GPU(duration=20)
def process_img(
    image, 
    prompt,
    negative_prompt,
    cfg_scale, # <-- RE-ENABLED
    steps, 
    img_strength, 
    depth_strength, 
    edge_strength,
    seed
):
    if image is None:
        raise gr.Error("Please upload an image first.")
    
    try:
        print("--- Starting Generation ---")
        result = gen.predict(
            image, 
            prompt,
            negative_prompt=negative_prompt,
            guidance_scale=cfg_scale, # <-- RE-ENABLED
            num_inference_steps=steps,
            img2img_strength=img_strength,
            depth_strength=depth_strength,
            lineart_strength=edge_strength,
            seed=seed
        )
        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 Style", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎮 Face to Style
        Upload any image. If there is a face, we'll keep the identity. If not, we'll stylize 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=f"The trigger words '{Config.STYLE_TRIGGER}' are added automatically."
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt (Optional)", 
                placeholder="e.g., blurry, text, watermark, bad art...",
                value=Config.DEFAULT_NEGATIVE_PROMPT
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                seed = gr.Number(
                    label="Seed", 
                    value=-1, 
                    info="-1 for random", 
                    precision=0
                )
                
                # --- RE-ENABLED CFG/GUIDANCE SLIDER ---
                cfg_scale = gr.Slider(
                    elem_id="cfg_scale",
                    minimum=1.0, 
                    maximum=10.0, # Range for TCD+Style
                    step=0.1, 
                    value=Config.CGF_SCALE, # Default 4.0
                    label="Style Strength (Guidance)"
                )
                
                steps = gr.Slider(
                    elem_id="steps",
                    minimum=1, 
                    maximum=20, 
                    step=1, 
                    value=8, # TCD default
                    label="Steps Number"
                )
                img_strength = gr.Slider(
                    elem_id="img_strength",
                    minimum=0.1, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.IMG_STRENGTH, 
                    label="Image Strength (Img2Img)"
                )
                depth_strength = gr.Slider(
                    elem_id="depth_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.DEPTH_STRENGTH, 
                    label="DepthMap Strength"
                )
                edge_strength = gr.Slider(
                    elem_id="edge_strength",
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.EDGE_STRENGTH, 
                    label="EdgeMap Strength (LineArt)"
                )
            
            run_btn = gr.Button("Generate", variant="primary")
            
        with gr.Column(scale=1):
            output_img = gr.Image(label="Styled Result")
                
    # Event Handler
    all_inputs = [
        input_img, 
        prompt, 
        negative_prompt,
        cfg_scale, # <-- RE-ENABLED
        steps, 
        img_strength, 
        depth_strength, 
        edge_strength,
        seed
    ]
    
    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,
        show_api=True 
    )