import torch
import os

auth_token = os.getenv("HF_TOKEN")

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else None

from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, use_auth_token=auth_token, revision="fp16", torch_dtype=torch_dtype
).to(device)

def predict(prompt):
    return pipe(prompt).images[0]
    
import gradio as gr

gradio_ui = gr.Interface(
    fn=predict,
    title="Stable Diffusion Demo",
    description="Enter a description of an image you'd like to generate!",
    inputs=[
        gr.Textbox(lines=2, label="Paste some text here"),
    ],
    outputs=["image"],
    examples=[["a photograph of an astronaut riding a horse"]],
)

gradio_ui.launch()