File size: 1,265 Bytes
d44e208
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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()