Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler | |
| MODEL_ID = "timbrooks/instruct-pix2pix" | |
| pipe = None | |
| def load_pipe(): | |
| global pipe | |
| if pipe is not None: | |
| return pipe | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.float16 if device == "cuda" else torch.float32 | |
| p = StableDiffusionInstructPix2PixPipeline.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=dtype, | |
| safety_checker=None, # keeps it simple for demo; you can re-add later | |
| ) | |
| # Scheduler that often looks better for edits | |
| p.scheduler = EulerAncestralDiscreteScheduler.from_config(p.scheduler.config) | |
| p = p.to(device) | |
| pipe = p | |
| return pipe | |
| def edit_image(image: Image.Image, prompt: str, strength: float, guidance: float, steps: int, seed: int): | |
| if image is None: | |
| return None | |
| if not prompt or not prompt.strip(): | |
| return image | |
| p = load_pipe() | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| generator = torch.Generator(device=device) | |
| if seed >= 0: | |
| generator = generator.manual_seed(seed) | |
| # InstructPix2Pix expects an RGB PIL image | |
| image = image.convert("RGB") | |
| out = p( | |
| prompt=prompt, | |
| image=image, | |
| num_inference_steps=int(steps), | |
| image_guidance_scale=float(strength), # how much it follows the input image | |
| guidance_scale=float(guidance), # how much it follows the prompt | |
| generator=generator, | |
| ).images[0] | |
| return out | |
| with gr.Blocks(title="SNAP AI Editor") as demo: | |
| gr.Markdown("## SNAP AI Editor\nUpload an image and describe the edit you want.") | |
| with gr.Row(): | |
| input_img = gr.Image(type="pil", label="Input image") | |
| output_img = gr.Image(type="pil", label="Output image") | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| placeholder="Examples: 'put me in a tuxedo', 'remove acne and smooth skin', 'change hair to blonde'" | |
| ) | |
| with gr.Row(): | |
| strength = gr.Slider(0.5, 2.0, value=1.2, step=0.05, label="Keep Original (image_guidance)") | |
| guidance = gr.Slider(1.0, 12.0, value=7.0, step=0.5, label="Follow Prompt (guidance)") | |
| steps = gr.Slider(10, 40, value=25, step=1, label="Steps") | |
| seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)") | |
| btn = gr.Button("Submit") | |
| btn.click( | |
| fn=edit_image, | |
| inputs=[input_img, prompt, strength, guidance, steps, seed], | |
| outputs=[output_img], | |
| ) | |
| demo.queue().launch() | |