Update app.py

app.py

@@ -1,24 +1,34 @@
 import io
+import os
 import torch
 from fastapi import FastAPI
 from pydantic import BaseModel
 from diffusers import FluxPipeline
 import gradio as gr
 from PIL import Image
+from huggingface_hub import login
 
 # Initialize FastAPI
 app = FastAPI()
 
+# 1. Login using the Secret stored in the Space settings
+hf_token = os.getenv("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
+else:
+    print("Warning: HF_TOKEN not found in Secrets. Gated models may fail.")
+
 # Load Model Optimized for CPU
-# NOTE: "black-forest-labs/FLUX.1-schnell" is huge. 
-# For HF Free Tier, consider a quantized version like "sayakpaul/flux.1-schnell-8bit"
 model_id = "black-forest-labs/FLUX.1-schnell" 
 
+# Using float32 or bfloat16 for CPU compatibility
 pipe = FluxPipeline.from_pretrained(
     model_id, 
-    torch_dtype=torch.bfloat16
+    torch_dtype=torch.bfloat16,
+    use_auth_token=True
 )
-# Vital for Free Tier: Moves parts of the model to CPU/Disk as needed
+
+# Enable CPU offloading to stay within the ~16GB RAM limit
 pipe.enable_model_cpu_offload() 
 
 class PromptRequest(BaseModel):
@@ -26,15 +36,16 @@ class PromptRequest(BaseModel):
 
 @app.post("/generate")
 def generate_api(request: PromptRequest):
+    # num_inference_steps=4 is the sweet spot for Schnell
     image = pipe(
         request.prompt, 
-        num_inference_steps=4, # Schnell is optimized for 4 steps
+        num_inference_steps=4, 
         guidance_scale=0.0
     ).images[0]
     
     img_byte_arr = io.BytesIO()
     image.save(img_byte_arr, format='PNG')
-    return {"image": img_byte_arr.getvalue().hex()} # Or return as StreamingResponse
+    return {"image": img_byte_arr.getvalue().hex()}
 
 def gradio_generate(prompt):
     return pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
@@ -43,14 +54,15 @@ def gradio_generate(prompt):
 with gr.Blocks() as demo:
     gr.Markdown("# FLUX.1 [schnell] CPU Explorer")
     with gr.Row():
-        input_text = gr.Textbox(label="Enter Prompt")
+        input_text = gr.Textbox(label="Enter Prompt", placeholder="A futuristic city in the style of cyberpunk...")
         output_img = gr.Image(label="Generated Image")
     btn = gr.Button("Generate")
     btn.click(fn=gradio_generate, inputs=input_text, outputs=output_img)
 
-# Mount FastAPI into Gradio for Hugging Face compatibility
+# Mount FastAPI into Gradio
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
     import uvicorn
+    # Port 7860 is required for Hugging Face Spaces
     uvicorn.run(app, host="0.0.0.0", port=7860)