Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from diffusers import StableDiffusionInpaintPipeline, LCMScheduler | |
| from PIL import Image | |
| # 1. Load the most compatible inpainting model | |
| model_id = "runwayml/stable-diffusion-inpainting" | |
| print("Loading Stable Diffusion Inpainting to CPU...") | |
| pipe = StableDiffusionInpaintPipeline.from_pretrained( | |
| model_id, | |
| torch_dtype=torch.float32, | |
| low_cpu_mem_usage=True | |
| ) | |
| # 2. THE "OFF SWITCH" FOR NSFW | |
| # This bypasses the built-in safety checker completely. | |
| pipe.safety_checker = None | |
| pipe.requires_safety_checker = False | |
| # 3. CPU OPTIMIZATIONS | |
| pipe.to("cpu") | |
| pipe.enable_attention_slicing() | |
| def predict(image_data, prompt): | |
| if image_data is None or not prompt: | |
| return None | |
| # 1. Prep Images | |
| raw_bg = image_data["background"].convert("RGB") | |
| raw_layer = image_data["layers"][-1] | |
| raw_mask = raw_layer.split()[-1] | |
| raw_mask = raw_mask.point(lambda x: 255 if x > 0 else 0) | |
| orig_w, orig_h = raw_bg.size | |
| scale = 512 / max(orig_w, orig_h) | |
| new_w, new_h = int((orig_w * scale) // 8) * 8, int((orig_h * scale) // 8) * 8 | |
| base_image = raw_bg.resize((new_w, new_h), Image.LANCZOS) | |
| mask_image = raw_mask.resize((new_w, new_h), Image.NEAREST) | |
| # return mask_image.resize((orig_w, orig_h), Image.NEAREST) | |
| # 3. RUN THE MODEL | |
| result = pipe( | |
| prompt=prompt, | |
| image=base_image, # Use the noised version! | |
| mask_image=mask_image, | |
| num_inference_steps=20, # Increased slightly for better detail on CPU | |
| guidance_scale=8.0, # Increased to make the "Cat" more likely to appear | |
| ).images[0] | |
| # 4. Final Alignment & Composite | |
| result = result.resize(base_image.size, Image.LANCZOS) | |
| final_segmented = Image.composite(result, base_image, mask_image) | |
| return final_segmented.resize((orig_w, orig_h), Image.LANCZOS) | |
| # 4. UI SETUP | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# 🎨 Fast CPU Inpainter (Unrestricted)") | |
| gr.Markdown("Draw on the image and describe what to add. **CPU Time: ~45 seconds.**") | |
| with gr.Row(): | |
| input_img = gr.ImageMask(label="Upload & Mask", type="pil") | |
| output_img = gr.Image(label="Result") | |
| prompt_text = gr.Textbox(label="Prompt") | |
| run_btn = gr.Button("Generate", variant="primary") | |
| run_btn.click(predict, [input_img, prompt_text], output_img) | |
| demo.launch() | |