File size: 1,642 Bytes
d447be3
096155b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1a0f0c4
096155b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image, make_image_grid
from PIL import Image
import requests
from io import BytesIO

# Load the pipeline
pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0", 
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True
)

# Offload model to reduce memory usage
pipeline.enable_model_cpu_offload()

# Function to load the initial image from a URL
def load_init_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content))

# Gradio function for image generation
def generate_image(prompt, image_url, strength):
    init_image = load_init_image(image_url)  # Load the initial image from the URL
    result_image = pipeline(prompt, image=init_image, strength=strength).images[0]
    
    # Display both the initial and result images side by side
    grid_image = make_image_grid([init_image, result_image], rows=1, cols=2)
    return grid_image

# Gradio interface
gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(lines=1, label="Prompt", placeholder="Enter the image description prompt"),
        gr.Textbox(lines=1, label="Image URL", placeholder="Enter the URL of the initial image"),
        gr.Slider(0.0, 1.0, value=0.5, label="Strength"),
    ],
    outputs=gr.Image(label="Image Comparison"),
    title="Stable Diffusion XL Refiner - Image to Image",
    description="Generate an image transformation from an initial image and a text prompt using the Stable Diffusion XL Refiner model.",
).launch()