import torch
from diffusers import StableDiffusionPipeline
import gradio as gr
import os
import spaces

HF_TOKEN = os.getenv['HF_TOKEN']

if torch.cuda.is_available():
    device = "cuda"
    print("Using GPU")
else:
    device = "cpu"
    print("Using CPU")

# Initialize the pipeline and download the model
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium", torch_dtype=torch.float16, use_auth_token=HF_TOKEN)
pipe.to(device)

# Define the image generation function
@spaces.GPU(duration=60)
def generate_image(prompt):
    image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
    return image

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(placeholder="Insert Prompt Here..."),
    outputs="image",
    title="Stable Diffusion 3 Medium",
    description="Made by [Nick088](https://linktr.ee/Nick088"
)

# Launch the interface
interface.launch(share=False)