clockclock commited on
Commit
77c7489
·
verified ·
1 Parent(s): 173fd35

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -31
app.py CHANGED
@@ -8,34 +8,29 @@ import time
8
 
9
  # --- Model Loading (CPU Version) ---
10
  print("Loading model on CPU... This may take several minutes.")
 
11
  pipe = AutoPipelineForInpainting.from_pretrained(
12
- "runwayml/stable-diffusion-inpainting" # Using the slightly smaller 1.5 model for better CPU performance
 
 
13
  )
14
  print("Model loaded successfully.")
15
 
 
16
  # --- Default "Magic" Prompts ---
17
- # These will be used if the user doesn't provide their own prompt.
18
  DEFAULT_PROMPT = "photorealistic, 4k, ultra high quality, sharp focus, masterpiece, high detail, professional photo"
19
  DEFAULT_NEGATIVE_PROMPT = "blurry, pixelated, distorted, deformed, ugly, disfigured, cartoon, anime, low quality, watermark, text"
20
 
21
  # --- The Inpainting Function ---
22
- # It now handles the logic for an optional user prompt.
23
  def inpaint_image(input_dict, user_prompt, guidance_scale, num_steps, progress=gr.Progress()):
24
- """
25
- Performs inpainting. Uses a default prompt if the user_prompt is empty.
26
- """
27
  image = input_dict["image"].convert("RGB")
28
  mask_image = input_dict["mask"].convert("RGB")
29
 
30
- # --- This is the core logic for the hybrid approach ---
31
  if user_prompt and user_prompt.strip():
32
- # If the user provided a prompt, use it.
33
  prompt = user_prompt
34
- # For custom prompts, a general negative prompt is still useful.
35
  negative_prompt = DEFAULT_NEGATIVE_PROMPT
36
  print(f"Using custom prompt: '{prompt}'")
37
  else:
38
- # If the user left the prompt box empty, use our high-quality defaults.
39
  prompt = DEFAULT_PROMPT
40
  negative_prompt = DEFAULT_NEGATIVE_PROMPT
41
  print(f"User prompt is empty. Using default 'General Fix' prompt.")
@@ -43,11 +38,9 @@ def inpaint_image(input_dict, user_prompt, guidance_scale, num_steps, progress=g
43
  print(f"Starting inpainting on CPU...")
44
  start_time = time.time()
45
 
46
- # Callback to update the progress bar in the UI
47
  def progress_callback(step, timestep, latents):
48
  progress(step / int(num_steps), desc=f"Running step {step}/{int(num_steps)}")
49
 
50
- # Run the pipeline
51
  result_image = pipe(
52
  prompt=prompt,
53
  image=image,
@@ -61,7 +54,6 @@ def inpaint_image(input_dict, user_prompt, guidance_scale, num_steps, progress=g
61
 
62
  end_time = time.time()
63
  print(f"Inpainting finished in {end_time - start_time:.2f} seconds.")
64
-
65
  return result_image
66
 
67
 
@@ -70,7 +62,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
70
  gr.Markdown(
71
  """
72
  # 🎨 AI Image Fixer
73
-
74
  **How to use:**
75
  1. Upload an image.
76
  2. Use the brush to **paint over the area you want to fix**.
@@ -82,31 +73,17 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
82
 
83
  gr.Warning(
84
  "⚠️ This Space is running on a free CPU. "
85
- "Image generation will be VERY SLOW (expect 5-15 minutes). "
86
  "A progress bar will show the status below the button. Please be patient!"
87
  )
88
 
89
  with gr.Row():
90
- # Input Column
91
  with gr.Column(scale=2):
92
- input_image = gr.Image(
93
- label="1. Upload & Mask Image",
94
- source="upload",
95
- tool="brush",
96
- type="pil"
97
- )
98
-
99
- # The prompt textbox is back, but now it's optional!
100
- prompt_textbox = gr.Textbox(
101
- label="2. Describe Your Fix (Optional)",
102
- placeholder="Leave empty for a general fix, or type e.g., 'a perfect human hand'"
103
- )
104
-
105
  with gr.Accordion("Advanced Settings", open=False):
106
  guidance_scale = gr.Slider(minimum=0, maximum=20, value=8.0, label="Guidance Scale")
107
  num_steps = gr.Slider(minimum=10, maximum=50, step=1, value=25, label="Inference Steps")
108
-
109
- # Output Column
110
  with gr.Column(scale=1):
111
  output_image = gr.Image(label="Result", type="pil")
112
 
 
8
 
9
  # --- Model Loading (CPU Version) ---
10
  print("Loading model on CPU... This may take several minutes.")
11
+ # This is the corrected line that fixes the error.
12
  pipe = AutoPipelineForInpainting.from_pretrained(
13
+ "stabilityai/stable-diffusion-2-inpainting",
14
+ torch_dtype=torch.float32, # Use float32 for broad CPU compatibility
15
+ safety_checker=None # Explicitly disable the safety checker to prevent loading errors
16
  )
17
  print("Model loaded successfully.")
18
 
19
+
20
  # --- Default "Magic" Prompts ---
 
21
  DEFAULT_PROMPT = "photorealistic, 4k, ultra high quality, sharp focus, masterpiece, high detail, professional photo"
22
  DEFAULT_NEGATIVE_PROMPT = "blurry, pixelated, distorted, deformed, ugly, disfigured, cartoon, anime, low quality, watermark, text"
23
 
24
  # --- The Inpainting Function ---
 
25
  def inpaint_image(input_dict, user_prompt, guidance_scale, num_steps, progress=gr.Progress()):
 
 
 
26
  image = input_dict["image"].convert("RGB")
27
  mask_image = input_dict["mask"].convert("RGB")
28
 
 
29
  if user_prompt and user_prompt.strip():
 
30
  prompt = user_prompt
 
31
  negative_prompt = DEFAULT_NEGATIVE_PROMPT
32
  print(f"Using custom prompt: '{prompt}'")
33
  else:
 
34
  prompt = DEFAULT_PROMPT
35
  negative_prompt = DEFAULT_NEGATIVE_PROMPT
36
  print(f"User prompt is empty. Using default 'General Fix' prompt.")
 
38
  print(f"Starting inpainting on CPU...")
39
  start_time = time.time()
40
 
 
41
  def progress_callback(step, timestep, latents):
42
  progress(step / int(num_steps), desc=f"Running step {step}/{int(num_steps)}")
43
 
 
44
  result_image = pipe(
45
  prompt=prompt,
46
  image=image,
 
54
 
55
  end_time = time.time()
56
  print(f"Inpainting finished in {end_time - start_time:.2f} seconds.")
 
57
  return result_image
58
 
59
 
 
62
  gr.Markdown(
63
  """
64
  # 🎨 AI Image Fixer
 
65
  **How to use:**
66
  1. Upload an image.
67
  2. Use the brush to **paint over the area you want to fix**.
 
73
 
74
  gr.Warning(
75
  "⚠️ This Space is running on a free CPU. "
76
+ "Image generation will be VERY SLOW (expect 5-20 minutes). "
77
  "A progress bar will show the status below the button. Please be patient!"
78
  )
79
 
80
  with gr.Row():
 
81
  with gr.Column(scale=2):
82
+ input_image = gr.Image(label="1. Upload & Mask Image", source="upload", tool="brush", type="pil")
83
+ 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'")
 
 
 
 
 
 
 
 
 
 
 
84
  with gr.Accordion("Advanced Settings", open=False):
85
  guidance_scale = gr.Slider(minimum=0, maximum=20, value=8.0, label="Guidance Scale")
86
  num_steps = gr.Slider(minimum=10, maximum=50, step=1, value=25, label="Inference Steps")
 
 
87
  with gr.Column(scale=1):
88
  output_image = gr.Image(label="Result", type="pil")
89