import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from PIL import Image
import requests
from io import BytesIO

# Set up the Stable Diffusion pipeline for text to image
text_to_image_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
text_to_image_pipe = text_to_image_pipe.to("cuda")

# Set up the Stable Diffusion pipeline for image to image
image_to_image_pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
device = "cuda"
image_to_image_pipe = image_to_image_pipe.to(device)

def text_to_image(prompt):
    image = text_to_image_pipe(prompt).images[0]
    image = Image.fromarray(image)
    return image

def image_to_image(file, prompt):
    response = requests.get(file)
    init_image = Image.open(BytesIO(response.content)).convert("RGB")
    init_image = init_image.resize((768, 512))
    images = image_to_image_pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
    return images[0]

# Create the Gradio interface
image_to_image_interface = gr.inputs.Image(label="Input Image")
text_prompt_interface = gr.inputs.Textbox(label="Text Prompt")

io_choice = gr.dropdown(["Text to Image", "Image to Image"], label="Choose Input Type")

outputs = []

if io_choice.lower() == "text to image":
    outputs.append(gr.outputs.Image(label="Generated Image"))
    interface = gr.Interface(fn=text_to_image, inputs=text_prompt_interface, outputs=outputs)
else:
    outputs.append(gr.outputs.Image(label="Generated Image"))
    interface = gr.Interface(fn=image_to_image, inputs=[image_to_image_interface, text_prompt_interface], outputs=outputs)

# Launch the interface
interface.launch()