app.py

# app.py (Final Corrected Version)

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

# --- Model Loading ---
print("Loading model for low-RAM CPU environment...")
model_id = "runwayml/stable-diffusion-inpainting"
try:
    pipe = AutoPipelineForInpainting.from_pretrained(model_id, torch_dtype=torch.float32)
    pipe.enable_model_cpu_offload()
    print("Model loaded successfully with CPU offloading enabled.")
except Exception as e:
    print(f"An error occurred during model loading: {e}")
    raise e

# --- Prompts ---
DEFAULT_PROMPT = "photorealistic, 4k, ultra high quality, sharp focus, masterpiece, high detail"
DEFAULT_NEGATIVE_PROMPT = "blurry, pixelated, distorted, deformed, ugly, disfigured, cartoon, watermark"

# --- Inpainting Function (Correct Signature) ---
def inpaint_image(image_and_mask, user_prompt, guidance_scale, num_steps, progress=gr.Progress(track_tqdm=True)):
    # The input is now a dictionary with 'image' and 'mask' keys
    image = image_and_mask["image"].convert("RGB")
    mask = image_and_mask["mask"].convert("RGB")
    
    if image is None or mask is None:
        raise gr.Error("Please upload an image and draw a mask on it first!")
    
    if user_prompt and user_prompt.strip():
        prompt = user_prompt
        negative_prompt = DEFAULT_NEGATIVE_PROMPT 
    else:
        prompt = DEFAULT_PROMPT
        negative_prompt = DEFAULT_NEGATIVE_PROMPT
    
    print(f"Starting inpainting on CPU...")
    result_image = pipe(
        prompt=prompt, image=image, mask_image=mask, negative_prompt=negative_prompt,
        guidance_scale=guidance_scale, num_inference_steps=int(num_steps)
    ).images[0]
    
    return result_image

# --- UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 AI Image Fixer (Stable Version)")
    gr.Warning("‼️ PATIENCE REQUIRED! Generation can take 15-30 minutes on free hardware.")
    
    with gr.Row():
        with gr.Column(scale=2):
            # This component returns a dictionary when tool='brush'
            input_image = gr.Image(label="1. Upload & Mask Image", source="upload", tool="brush", type="pil")
            prompt_textbox = gr.Textbox(label="2. Describe Your Fix (Optional)", placeholder="Leave empty for a general fix")
            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=20, label="Inference Steps")
        with gr.Column(scale=1):
            output_image = gr.Image(label="Result", type="pil")

    submit_button = gr.Button("Fix It!", variant="primary")
    
    # The `inputs` list is simple. The function signature must match what Gradio provides.
    submit_button.click(
        fn=inpaint_image,
        inputs=[input_image, prompt_textbox, guidance_scale, num_steps],
        outputs=output_image
    )

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