import gradio as gr
from PIL import Image
import torch
from torch import autocast
from diffusers import StableDiffusionImg2ImgPipeline
from torch import autocast
from tqdm.auto import tqdm
import requests
from io import BytesIO
from PIL import Image
from typing import List, Optional, Union
import inspect
import warnings

# Load the Stable Diffusion model
modelid = "CompVis/stable-diffusion-v1-4"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(modelid, revision="fp16", torch_dtype=torch.float16)
pipe.to(device)

def generate_image(prompt):
    response = requests.get(url)
    init_img = Image.open(BytesIO(response.content)).convert("RGB")
    init_img = init_img.resize((768, 512))

    generator = torch.Generator(device=device).manual_seed(1024)
    with autocast("cuda"):
        image = pipe(prompt=prompt, init_image=init_img, strength=0.75, guidance_scale=7.5, generator=generator).images[0]

    return image

# Define the input and output components
input_text = gr.inputs.Textbox(lines=10, label="Enter a prompt")
output_image = gr.outputs.Image(type="pil", label="Generated Image")

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=input_text,
    outputs=output_image,
    title="Stable Bud",
    description="Generate images using Stable Diffusion",
    layout="vertical",
)

if __name__ == "__main__":
    iface.launch()