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()