Update app.py

app.py

@@ -1,36 +1,53 @@
-# app.py (Modified for CPU)
+# app.py
 
 import gradio as gr
 import torch
 from diffusers import AutoPipelineForInpainting
 from PIL import Image
-import time # Import time to measure execution
+import time
 
 # --- 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.")
+print("Loading model on CPU... This may take several minutes.")
 pipe = AutoPipelineForInpainting.from_pretrained(
-    "stabilityai/stable-diffusion-2-inpainting"
+    "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 ---
-def inpaint_image(input_dict, prompt, negative_prompt, guidance_scale, num_steps, progress=gr.Progress()):
+# 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 on an image based on a mask and a prompt.
-    Includes progress tracking for the slow CPU process.
+    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")
 
-    print(f"Starting inpainting with prompt: '{prompt}' on CPU.")
+    # --- 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,
@@ -47,43 +64,57 @@ def inpaint_image(input_dict, prompt, negative_prompt, guidance_scale, num_steps
 
     return result_image
 
+
 # --- Gradio User Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🎨 AI Image Fixer (Inpainting)
-        Upload an image, mask the area to fix, and describe the change.
+        # 🎨 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!"
         """
     )
     
-    # *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."
+        "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():
-        with gr.Column():
+        # Input Column
+        with gr.Column(scale=2):
             input_image = gr.Image(
-                label="Upload Image & Draw Mask", source="upload", tool="brush", type="pil"
+                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'"
             )
-            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'")
-                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")
+                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")
         
-        with gr.Column():
-            output_image = gr.Image(label="Resulting Image", type="pil")
+        # Output Column
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Result", type="pil")
 
     submit_button = gr.Button("Fix It!", variant="primary")
     
-    # We add a progress component to be updated
     submit_button.click(
         fn=inpaint_image,
-        inputs=[input_image, prompt, negative_prompt, guidance_scale, num_steps],
+        inputs=[input_image, prompt_textbox, guidance_scale, num_steps],
         outputs=output_image
     )