Spaces:
Sleeping
Sleeping
| import torch | |
| from diffusers import StableDiffusionInstructPix2PixPipeline | |
| import gradio as gr | |
| from PIL import Image | |
| # Ensure we use GPU if available | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def load_model(): | |
| """Load InstructPix2Pix model.""" | |
| pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained( | |
| "timbrooks/instruct-pix2pix", | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32 | |
| ).to(device) | |
| return pipe | |
| pipe = load_model() | |
| def edit_image(image, prompt): | |
| """Perform image editing using InstructPix2Pix model.""" | |
| if image is None: | |
| raise ValueError("Please upload an image.") | |
| image = Image.fromarray(image).convert("RGB") # Ensure image is in RGB format | |
| result = pipe(image=image, prompt=prompt, num_inference_steps=10, image_guidance_scale=1.5).images[0] | |
| return result | |
| # Gradio UI | |
| demo = gr.Interface( | |
| fn=edit_image, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload Image"), | |
| gr.Textbox(label="Editing Instruction"), | |
| ], | |
| outputs=gr.Image(type="pil", label="Edited Image"), | |
| title="InstructPix2Pix Image Editing", | |
| description="Upload an image and provide a text instruction to modify it.", | |
| ) | |
| demo.launch() | |