Spaces:
Sleeping
Sleeping
File size: 3,791 Bytes
173fd35 55d13a0 173fd35 55d13a0 c50ee85 173fd35 77c7489 c50ee85 77c7489 c50ee85 55d13a0 77c7489 173fd35 55d13a0 173fd35 55d13a0 173fd35 c50ee85 55d13a0 c50ee85 55d13a0 c50ee85 55d13a0 c50ee85 55d13a0 173fd35 55d13a0 173fd35 55d13a0 c50ee85 77c7489 173fd35 c50ee85 55d13a0 173fd35 77c7489 55d13a0 173fd35 55d13a0 173fd35 55d13a0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
# app.py
import gradio as gr
import torch
from diffusers import AutoPipelineForInpainting
from PIL import Image
import time
# --- Model Loading (CPU Version) ---
print("Loading model on CPU... This may take several minutes.")
# This is the corrected line that fixes the error.
pipe = AutoPipelineForInpainting.from_pretrained(
"stabilityai/stable-diffusion-2-inpainting",
torch_dtype=torch.float32, # Use float32 for broad CPU compatibility
safety_checker=None # Explicitly disable the safety checker to prevent loading errors
)
print("Model loaded successfully.")
# --- 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 ---
def inpaint_image(input_dict, user_prompt, guidance_scale, num_steps, progress=gr.Progress()):
image = input_dict["image"].convert("RGB")
mask_image = input_dict["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...")
start_time = time.time()
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,
mask_image=mask_image,
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
# --- Gradio User Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# 🎨 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!"
"""
)
gr.Warning(
"⚠️ This Space is running on a free CPU. "
"Image generation will be VERY SLOW (expect 5-20 minutes). "
"A progress bar will show the status below the button. Please be patient!"
)
with gr.Row():
with gr.Column(scale=2):
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, or type e.g., 'a perfect human hand'")
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=25, label="Inference Steps")
with gr.Column(scale=1):
output_image = gr.Image(label="Result", type="pil")
submit_button = gr.Button("Fix It!", variant="primary")
submit_button.click(
fn=inpaint_image,
inputs=[input_image, prompt_textbox, guidance_scale, num_steps],
outputs=output_image
)
if __name__ == "__main__":
demo.launch() |