Update app.py

app.py

@@ -1,27 +1,32 @@
+import os
 import gradio as gr
 import torch
 import spaces
 from diffusers import FluxPipeline
-import os
+from huggingface_hub import login
+
+# --- AUTHENTICATION LAYER ---
+# This ensures the session is logged in BEFORE loading the model
+token = os.getenv("FLUX_API_TOKEN")
+if token:
+    login(token=token)
+else:
+    print("❌ Error: HF_TOKEN secret not found. Go to Settings > Secrets.")
 
-# 1. Initialize the Pipeline
-hf_token = os.getenv("FLUX_API_TOKEN")
+# --- MODEL LOADING ---
 model_id = "black-forest-labs/FLUX.1-schnell"
 
-# Load the model with bfloat16 to save memory and improve speed
+# We pass the token explicitly here as well
 pipe = FluxPipeline.from_pretrained(
     model_id, 
     torch_dtype=torch.bfloat16,
-    hf_token = os.getenv("FLUX_API_TOKEN")
+    token=token
 )
 
-# 2. Define the Generation Function with ZeroGPU decorator
-# The @spaces.GPU decorator handles the dynamic GPU allocation on Hugging Face
 @spaces.GPU(duration=60)
 def generate_image(prompt, seed, width, height, steps):
-    pipe.to("cuda") # Moves model to GPU only during execution
-    
-    generator = torch.Generator("cuda").manual_seed(seed)
+    pipe.to("cuda")
+    generator = torch.Generator("cuda").manual_seed(int(seed))
     
     image = pipe(
         prompt=prompt,
@@ -29,36 +34,26 @@ def generate_image(prompt, seed, width, height, steps):
         height=height,
         num_inference_steps=steps, 
         generator=generator,
-        guidance_scale=0.0 # Schnell version works best with 0 guidance
+        guidance_scale=0.0
     ).images[0]
     
     return image
 
-# 3. Create the Gradio Interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 FLUX.1-schnell Image Generator")
-    gr.Markdown("Generating high-quality images on Hugging Face ZeroGPU.")
-
+# --- GRADIO INTERFACE ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 🚀 FLUX.1 Schnell - Direct Access")
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Enter your prompt", placeholder="A neon-lit cyberpunk cat...")
-            generate_btn = gr.Button("Generate", variant="primary")
-            
-            with gr.Accordion("Settings", open=False):
-                seed = gr.Slider(0, 1000000, label="Seed", value=42, step=1)
-                width = gr.Slider(512, 1024, label="Width", value=1024, step=64)
-                height = gr.Slider(512, 1024, label="Height", value=1024, step=64)
-                steps = gr.Slider(1, 4, label="Inference Steps", value=4, step=1)
-        
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image")
+            prompt = gr.Textbox(label="Prompt", placeholder="A glass sculpture of a bird...")
+            btn = gr.Button("Generate", variant="primary")
+            with gr.Accordion("Advanced", open=False):
+                seed = gr.Number(label="Seed", value=42)
+                width = gr.Slider(512, 1024, value=1024, step=64)
+                height = gr.Slider(512, 1024, value=1024, step=64)
+                steps = gr.Slider(1, 4, value=4, step=1)
+        output_img = gr.Image(label="Result")
 
-    generate_btn.click(
-        fn=generate_image, 
-        inputs=[prompt, seed, width, height, steps], 
-        outputs=output_image
-    )
+    btn.click(fn=generate_image, inputs=[prompt, seed, width, height, steps], outputs=output_img)
 
-# 4. Launch the App
 if __name__ == "__main__":
     demo.launch()