File size: 1,502 Bytes
d2adf52
7b9863b
d2adf52
 
7b9863b
 
68b1542
 
 
 
7b9863b
 
 
 
 
 
 
 
68b1542
 
7b9863b
68b1542
 
7b9863b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d2adf52
7b9863b
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
46
47
import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionPipeline
import torch

# Text-to-Image pipeline
txt2img_pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16
).to("cuda")

# Image-to-Image pipeline
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16
).to("cuda")

# Text-to-Image function
def text_to_image(prompt):
    if not prompt:
        return None
    image = txt2img_pipe(prompt).images[0]
    return image

# Image-to-Image function
def image_to_image(image, prompt):
    if image is None or not prompt:
        return None
    output = img2img_pipe(prompt=prompt, image=image, strength=0.75).images[0]
    return output

# Gradio Interface with Tabs
with gr.Blocks() as demo:
    with gr.Tab("Text → Image"):
        txt_prompt = gr.Textbox(label="Enter prompt")
        txt_output = gr.Image(label="Generated Image")
        txt_btn = gr.Button("Generate")
        txt_btn.click(fn=text_to_image, inputs=txt_prompt, outputs=txt_output)

    with gr.Tab("Image + Text → Image"):
        img_input = gr.Image(label="Upload image")
        img_prompt = gr.Textbox(label="Enter prompt")
        img_output = gr.Image(label="Modified Image")
        img_btn = gr.Button("Generate")
        img_btn.click(fn=image_to_image, inputs=[img_input, img_prompt], outputs=img_output)

demo.launch()