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

from huggingface_hub import login

# 3. Login to Hugging Face (replace YOUR_HF_TOKEN with your actual token)


login(token=os.getenv("HF_TOKEN"))

# 4. Load private model using your token
model_id = "Lookingsoft-team/Text_to_Image_Diffusion"
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16  # smaller memory
).to(device)

def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter prompt"),
    outputs=gr.Image(type="pil"),
    title="Private Stable Diffusion Demo"
)

iface.launch(share=True)