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

# Retrieve the token from the environment variable
token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
if token:
    login(token=token)  # Log in with the retrieved token
else:
    raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")

# Load the Stable Diffusion 3.5 model
model_id = "stabilityai/stable-diffusion-3.5-large"
pipe = StableDiffusion3Pipeline.from_pretrained(model_id)  # No LoRA integration

pipe.to("cpu")  # Ensuring it runs on CPU

# Function to generate an image from a text prompt
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio interface
iface = gr.Interface(fn=generate_image, inputs="text", outputs="image")
iface.launch()