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

# मॉडल आईडी
model_id = "black-forest-labs/FLUX.1-dev"

# Gated मॉडल के लिए टोकन और ऑथेंटिकेशन
hf_token = os.environ.get("HF_TOKEN")

# पाइपलाइन लोड करें (token पैरामीटर का उपयोग करके)
pipe = FluxPipeline.from_pretrained(
    model_id, 
    token=hf_token, 
    torch_dtype=torch.bfloat16
)
pipe = pipe.to("cuda")

@spaces.GPU
def generate_image(prompt):
    # ZeroGPU के लिए जरूरी @spaces.GPU डेकोरेटर
    image = pipe(
        prompt,
        height=768,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=20,
        max_sequence_length=512,
    ).images[0]
    return image

# Gradio इंटरफेस
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt", placeholder="Describe your realistic image..."),
    outputs=gr.Image(label="Generated Image"),
    title="Flux.1 Dev - Gated Model"
)

if __name__ == "__main__":
    demo.launch()