Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Final Corrected Version)
+# app.py
 
 import gradio as gr
 import torch
@@ -6,64 +6,95 @@ 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"
+# --- Model Loading (Final, Most Stable Version) ---
+print("Loading the definitive model for low-RAM CPU environment...")
+# We are using the more modern and reliable SD 2.0 Inpainting model.
+# This model is better packaged and less prone to loading errors.
+model_id = "stabilityai/stable-diffusion-2-inpainting"
+
 try:
-    pipe = AutoPipelineForInpainting.from_pretrained(model_id, torch_dtype=torch.float32)
+    pipe = AutoPipelineForInpainting.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32, # Use float32 for CPU compatibility
+        safety_checker=None        # Proactively disable the safety checker to save memory
+    )
+    
+    # Enable CPU offloading to prevent memory crashes. This is essential.
     pipe.enable_model_cpu_offload()
-    print("Model loaded successfully with CPU offloading enabled.")
+    
+    print("Model loaded successfully. The application is ready.")
+
 except Exception as e:
-    print(f"An error occurred during model loading: {e}")
+    print("="*80)
+    print("A FATAL ERROR OCCURRED DURING MODEL LOADING. The app cannot start.")
+    print(f"Error: {e}")
+    print("This is likely due to the free hardware tier not having enough resources.")
+    print("="*80)
     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) ---
+# --- Default "Magic" Prompts ---
+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 ---
+# This function signature is correct for how Gradio's Image tool works.
 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:
+    if image_and_mask is None or "image" not in image_and_mask or "mask" not in image_and_mask:
         raise gr.Error("Please upload an image and draw a mask on it first!")
     
+    # The input is a dictionary with 'image' and 'mask' keys
+    image = image_and_mask["image"].convert("RGB")
+    mask = image_and_mask["mask"].convert("RGB")
+    
     if user_prompt and user_prompt.strip():
         prompt = user_prompt
         negative_prompt = DEFAULT_NEGATIVE_PROMPT 
+        print(f"Using custom prompt: '{prompt}'")
     else:
         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 (with offloading)... This will be very slow.")
+    start_time = time.time()
     
-    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)
+        prompt=prompt,
+        image=image,
+        mask_image=mask,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=int(num_steps),
     ).images[0]
     
+    end_time = time.time()
+    print(f"Inpainting finished in {end_time - start_time:.2f} seconds.")
     return result_image
 
-# --- UI ---
+
+# --- Gradio User Interface ---
 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.")
+    gr.Markdown("# 🎨 AI Image Fixer (Definitive Version)")
+    gr.Warning(
+        "‼️ **PATIENCE REQUIRED!** This app is running on a free CPU. "
+        "Generation will be **extremely slow (potentially 20-40 minutes)** due to memory-saving measures. "
+        "This is necessary to prevent crashes. The progress bar will appear after you click the button."
+    )
     
     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")
+                num_steps = gr.Slider(minimum=10, maximum=50, step=1, value=20, label="Inference Steps (Fewer is faster)")
         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],