Update app.py

app.py

@@ -1,53 +1,36 @@
-# app.py
+# app.py (Modified for CPU)
 
 import gradio as gr
 import torch
 from diffusers import AutoPipelineForInpainting
 from PIL import Image
+import time # Import time to measure execution
 
-# --- Model Loading ---
-# Load the model only once at the start of the application
-# We use float16 for memory efficiency and speed on GPUs
-# If no GPU is available, this will run on CPU (but it will be very slow)
-try:
-    pipe = AutoPipelineForInpainting.from_pretrained(
-        "stabilityai/stable-diffusion-2-inpainting",
-        torch_dtype=torch.float16,
-        variant="fp16"
-    ).to("cuda")
-except Exception as e:
-    print(f"Could not load model on GPU: {e}. Falling back to CPU.")
-    pipe = AutoPipelineForInpainting.from_pretrained(
-        "stabilityai/stable-diffusion-2-inpainting"
-    )
+# --- Model Loading (CPU Version) ---
+# We load the model without GPU-specific options.
+# This will run on the CPU.
+print("Loading model on CPU... This may take a moment.")
+pipe = AutoPipelineForInpainting.from_pretrained(
+    "stabilityai/stable-diffusion-2-inpainting"
+)
+print("Model loaded successfully.")
 
 # --- The Inpainting Function ---
-# This is the core function that takes user inputs and generates the image
-def inpaint_image(input_dict, prompt, negative_prompt, guidance_scale, num_steps):
+def inpaint_image(input_dict, prompt, negative_prompt, guidance_scale, num_steps, progress=gr.Progress()):
     """
     Performs inpainting on an image based on a mask and a prompt.
-    
-    Args:
-        input_dict (dict): A dictionary from Gradio's Image component containing 'image' and 'mask'.
-        prompt (str): The text prompt describing what to generate in the masked area.
-        negative_prompt (str): The text prompt describing what to avoid.
-        guidance_scale (float): A value to control how much the generation follows the prompt.
-        num_steps (int): The number of inference steps.
-        
-    Returns:
-        PIL.Image: The resulting image after inpainting.
+    Includes progress tracking for the slow CPU process.
     """
-    # Separate the image and the mask from the input dictionary
     image = input_dict["image"].convert("RGB")
     mask_image = input_dict["mask"].convert("RGB")
 
-    # The model works best with images of a specific size (e.g., 512x512)
-    # We can resize for consistency, but for user-friendliness, we'll let the pipeline handle it.
-    # However, it's good practice to inform the user that square images work best.
+    print(f"Starting inpainting with prompt: '{prompt}' on CPU.")
+    start_time = time.time()
     
-    print(f"Starting inpainting with prompt: '{prompt}'")
-    
-    # Run the inpainting pipeline
+    # 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)}")
+
     result_image = pipe(
         prompt=prompt,
         image=image,
@@ -55,7 +38,12 @@ def inpaint_image(input_dict, prompt, negative_prompt, guidance_scale, num_steps
         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
 
@@ -64,57 +52,40 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         # 🎨 AI Image Fixer (Inpainting)
-        
-        Have an AI-generated image with weird hands, faces, or artifacts? Fix it here!
-        
-        **How to use:**
-        1. Upload your image.
-        2. Use the brush tool to "paint" over the parts you want to replace. This is your mask.
-        3. Write a prompt describing what you want in the painted-over area.
-        4. Adjust the advanced settings if you want more control.
-        5. Click "Fix It!" and see the magic happen.
+        Upload an image, mask the area to fix, and describe the change.
         """
     )
     
+    # *CRUCIAL* Warning for CPU users
+    gr.Warning(
+        "⚠️ This Space is running on a free CPU. "
+        "Image generation will be VERY SLOW (expect 5-15 minutes per image). "
+        "Please be patient! A progress bar will appear below the 'Fix It!' button."
+    )
+    
     with gr.Row():
-        # Input column
         with gr.Column():
-            gr.Markdown("### 1. Upload & Mask Your Image")
-            # The Image component with a drawing tool for masking
             input_image = gr.Image(
-                label="Upload Image & Draw Mask",
-                source="upload",
-                tool="brush",
-                type="pil" # We want to work with PIL images in our function
+                label="Upload Image & Draw Mask", source="upload", tool="brush", type="pil"
             )
-            
-            gr.Markdown("### 2. Describe Your Fix")
-            prompt = gr.Textbox(label="Prompt", placeholder="e.g., 'A beautiful, realistic human hand, detailed fingers'")
-            
-            # Accordion for advanced settings to keep the UI clean
+            prompt = gr.Textbox(label="Prompt", placeholder="e.g., 'A beautiful, realistic human hand'")
             with gr.Accordion("Advanced Settings", open=False):
-                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="e.g., 'blurry, distorted, extra fingers, cartoon'")
-                guidance_scale = gr.Slider(minimum=0, maximum=20, value=8.0, label="Guidance Scale")
-                num_steps = gr.Slider(minimum=10, maximum=100, step=1, value=40, label="Inference Steps")
+                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="e.g., 'blurry, distorted, extra fingers'")
+                guidance_scale = gr.Slider(minimum=0, maximum=20, value=7.5, label="Guidance Scale")
+                # Lower the default steps for faster (but lower quality) generation on CPU
+                num_steps = gr.Slider(minimum=5, maximum=50, step=1, value=20, label="Inference Steps")
         
-        # Output column
         with gr.Column():
-            gr.Markdown("### 3. Get Your Result")
-            output_image = gr.Image(
-                label="Resulting Image",
-                type="pil"
-            )
+            output_image = gr.Image(label="Resulting Image", type="pil")
 
-    # The button to trigger the process
     submit_button = gr.Button("Fix It!", variant="primary")
     
-    # Connect the button to the function
+    # We add a progress component to be updated
     submit_button.click(
         fn=inpaint_image,
         inputs=[input_image, prompt, negative_prompt, guidance_scale, num_steps],
         outputs=output_image
     )
 
-# Launch the Gradio app
 if __name__ == "__main__":
     demo.launch()