# app.py

import gradio as gr
import torch
from diffusers import AutoPipelineForInpainting
from PIL import Image
import time

# --- Model Loading (CPU Version) ---
print("Loading model on CPU... This may take several minutes.")
pipe = AutoPipelineForInpainting.from_pretrained(
    "runwayml/stable-diffusion-inpainting" # Using the slightly smaller 1.5 model for better CPU performance
)
print("Model loaded successfully.")

# --- Default "Magic" Prompts ---
# These will be used if the user doesn't provide their own prompt.
DEFAULT_PROMPT = "photorealistic, 4k, ultra high quality, sharp focus, masterpiece, high detail, professional photo"
DEFAULT_NEGATIVE_PROMPT = "blurry, pixelated, distorted, deformed, ugly, disfigured, cartoon, anime, low quality, watermark, text"

# --- The Inpainting Function ---
# It now handles the logic for an optional user prompt.
def inpaint_image(input_dict, user_prompt, guidance_scale, num_steps, progress=gr.Progress()):
    """
    Performs inpainting. Uses a default prompt if the user_prompt is empty.
    """
    image = input_dict["image"].convert("RGB")
    mask_image = input_dict["mask"].convert("RGB")

    # --- This is the core logic for the hybrid approach ---
    if user_prompt and user_prompt.strip():
        # If the user provided a prompt, use it.
        prompt = user_prompt
        # For custom prompts, a general negative prompt is still useful.
        negative_prompt = DEFAULT_NEGATIVE_PROMPT 
        print(f"Using custom prompt: '{prompt}'")
    else:
        # If the user left the prompt box empty, use our high-quality defaults.
        prompt = DEFAULT_PROMPT
        negative_prompt = DEFAULT_NEGATIVE_PROMPT
        print(f"User prompt is empty. Using default 'General Fix' prompt.")
    
    print(f"Starting inpainting on CPU...")
    start_time = time.time()
    
    # Callback to update the progress bar in the UI
    def progress_callback(step, timestep, latents):
        progress(step / int(num_steps), desc=f"Running step {step}/{int(num_steps)}")

    # Run the pipeline
    result_image = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask_image,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=int(num_steps),
        callback_steps=1,
        callback=progress_callback,
    ).images[0]
    
    end_time = time.time()
    print(f"Inpainting finished in {end_time - start_time:.2f} seconds.")

    return result_image


# --- Gradio User Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 AI Image Fixer
        
        **How to use:**
        1. Upload an image.
        2. Use the brush to **paint over the area you want to fix**.
        3. **(Optional)** For precise control, write a custom prompt describing the fix.
        4. **(Easy Mode)** Or, just leave the prompt box empty for a general quality improvement.
        5. Click "Fix It!"
        """
    )
    
    gr.Warning(
        "⚠️ This Space is running on a free CPU. "
        "Image generation will be VERY SLOW (expect 5-15 minutes). "
        "A progress bar will show the status below the button. Please be patient!"
    )
    
    with gr.Row():
        # Input Column
        with gr.Column(scale=2):
            input_image = gr.Image(
                label="1. Upload & Mask Image",
                source="upload",
                tool="brush",
                type="pil"
            )
            
            # The prompt textbox is back, but now it's optional!
            prompt_textbox = gr.Textbox(
                label="2. Describe Your Fix (Optional)",
                placeholder="Leave empty for a general fix, or type e.g., 'a perfect human hand'"
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                guidance_scale = gr.Slider(minimum=0, maximum=20, value=8.0, label="Guidance Scale")
                num_steps = gr.Slider(minimum=10, maximum=50, step=1, value=25, label="Inference Steps")
        
        # Output Column
        with gr.Column(scale=1):
            output_image = gr.Image(label="Result", type="pil")

    submit_button = gr.Button("Fix It!", variant="primary")
    
    submit_button.click(
        fn=inpaint_image,
        inputs=[input_image, prompt_textbox, guidance_scale, num_steps],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch()