import os 
from huggingface_hub import login
import torch
import gradio as gr
from diffusers import StableDiffusion3Pipeline

# Get Hugging Face token from environment variables
hf_token = os.getenv("HF_API_TOKEN")

if hf_token:
    login(token=hf_token)
    print("Login successful")
else:
    raise ValueError("Hugging Face token is missing. Please set it in the environment variables.")

def image_generator(prompt):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipeline = StableDiffusion3Pipeline.from_pretrained(
        "stabilityai/stable-diffusion-3-medium-diffusers", 
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        text_encoder_3=None,
        tokenizer_3=None
    )
    # Move the pipeline to the appropriate device
    pipeline.to(device)

    # Generate the image
    image = pipeline(
        prompt=prompt,
        negative_prompt="blurred, ugly, watermark, low, resolution, blurry",
        num_inference_steps=40,
        height=1024,
        width=1024,
        guidance_scale=9.0
    ).images[0]

    return image

# Create a Gradio interface
interface = gr.Interface(
    fn=image_generator, 
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."), 
    outputs=gr.Image(type="pil"),
    title="Image Generator App",
    description="This is a simple image generator app using HuggingFace's Stable Diffusion 3 model."
)

# Launch the interface
interface.launch()
print(interface)