import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import hf_api

# Retrieve the Hugging Face token stored in Hugging Face Spaces secrets
HUGGINGFACE_TOKEN = hf_api.get_secret("keyss")

if not HUGGINGFACE_TOKEN:
    raise ValueError("Hugging Face token not found! Make sure it's added in the Hugging Face Secrets.")

def image_generation(prompt):
    # Check if GPU is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Load the Stable Diffusion 3 pipeline
    pipeline = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-3-medium-diffusers",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        use_auth_token=HUGGINGFACE_TOKEN,  # Use the Hugging Face token for authentication
        text_encoder_3=None,
        tokenizer_3=None
    )
    
    # Enable efficient model execution
    pipeline.enable_model_cpu_offload()
    
    # Generate an image based on the prompt
    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

# Define the Gradio interface
interface = gr.Interface(
    fn=image_generation,
    inputs=gr.Textbox(lines=2, placeholder="Enter your Prompt..."),
    outputs=gr.Image(type="pil"),
    title="Image Creation using Stable Diffusion 3 Model",
    description="This application generates awesome images using the Stable Diffusion 3 model."
)

# Launch the Gradio app
interface.launch()