updating app.py with fixes

app.py

@@ -1,35 +1,57 @@
 import torch
-from diffusers import Flux2Pipeline, Flux2Transformer2DModel
-from diffusers.utils import load_image
+from diffusers import Flux2Pipeline
 from huggingface_hub import get_token
 import requests
 import io
 import gradio as gr
 from PIL import Image
+import os
 
+# Configuration
 repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
-device = "cuda:0"
-torch_dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-def remote_text_encoder(prompts):
-    response = requests.post(
-        "https://remote-text-encoder-flux-2.huggingface.co/predict",
-        json={"prompt": prompts},
-        headers={
-            "Authorization": f"Bearer {get_token()}",
-            "Content-Type": "application/json"
-        }
-    )
-    prompt_embeds = torch.load(io.BytesIO(response.content))
+print(f"Using device: {device}")
+print(f"Using dtype: {torch_dtype}")
 
-    return prompt_embeds.to(device)
+def remote_text_encoder(prompts):
+    """Encode prompts using remote text encoder API."""
+    try:
+        token = get_token()
+        if not token:
+            raise ValueError("HuggingFace token not found. Please login using 'huggingface-cli login'")
+        
+        response = requests.post(
+            "https://remote-text-encoder-flux-2.huggingface.co/predict",
+            json={"prompt": prompts},
+            headers={
+                "Authorization": f"Bearer {token}",
+                "Content-Type": "application/json"
+            },
+            timeout=60
+        )
+        response.raise_for_status()
+        prompt_embeds = torch.load(io.BytesIO(response.content))
+        return prompt_embeds.to(device)
+    except Exception as e:
+        raise Exception(f"Failed to encode prompt: {str(e)}")
 
 # Load the pipeline
 print("Loading Flux2 pipeline...")
-pipe = Flux2Pipeline.from_pretrained(
-    repo_id, text_encoder=None, torch_dtype=torch_dtype
-).to(device)
-print("Pipeline loaded successfully!")
+try:
+    pipe = Flux2Pipeline.from_pretrained(
+        repo_id,
+        text_encoder=None,
+        torch_dtype=torch_dtype,
+        device_map="cuda"
+    )
+    if not torch.cuda.is_available():
+        pipe = pipe.to(device)
+    print("Pipeline loaded successfully!")
+except Exception as e:
+    print(f"Error loading pipeline: {e}")
+    raise
 
 def generate_image(
     prompt: str,
@@ -52,6 +74,9 @@ def generate_image(
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter a prompt!")
     
+    if not torch.cuda.is_available():
+        raise gr.Error("This Space requires a GPU to run. Please try a different Space or upgrade your hardware.")
+    
     progress(0, desc="Encoding prompt...")
     
     try:
@@ -61,17 +86,19 @@ def generate_image(
         progress(0.3, desc="Generating image...")
         
         # Set up generator
+        generator_device = "cuda" if torch.cuda.is_available() else "cpu"
         if seed == -1:
-            generator = torch.Generator(device=device)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
+            import random
+            seed = random.randint(0, 2**32 - 1)
+        
+        generator = torch.Generator(device=generator_device).manual_seed(int(seed))
         
         # Prepare pipeline arguments
         pipe_kwargs = {
             "prompt_embeds": prompt_embeds,
             "generator": generator,
-            "num_inference_steps": num_inference_steps,
-            "guidance_scale": guidance_scale,
+            "num_inference_steps": int(num_inference_steps),
+            "guidance_scale": float(guidance_scale),
         }
         
         # Add input image if provided
@@ -80,22 +107,32 @@ def generate_image(
             progress(0.4, desc="Processing input image...")
         
         # Generate image
-        image = pipe(**pipe_kwargs).images[0]
+        with torch.inference_mode():
+            image = pipe(**pipe_kwargs).images[0]
         
         progress(1.0, desc="Done!")
         
         return image
     
     except Exception as e:
+        import traceback
+        error_msg = f"Error generating image: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
         raise gr.Error(f"Error generating image: {str(e)}")
 
+
 # Create Gradio interface
-with gr.Blocks(title="Flux2 Image Generator", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(
+    title="Flux2 Image Generator",
+) as demo:
     gr.Markdown(
         """
         # 🎨 Flux2 Image Generator
-        Generate stunning images using FLUX.2-dev with 4-bit quantization.
+        Generate stunning images using **FLUX.2-dev** with 4-bit quantization for efficient inference.
+        
         Supports both **text-to-image** and **image-to-image** generation.
+        
+        ⚠️ **Note**: This Space requires a GPU to run.
         """
     )
     
@@ -148,7 +185,8 @@ with gr.Blocks(title="Flux2 Image Generator", theme=gr.themes.Soft()) as demo:
             generate_btn = gr.Button(
                 "🚀 Generate Image",
                 variant="primary",
-                size="lg"
+                size="lg",
+                elem_classes="generate-btn"
             )
             
             gr.Markdown(
@@ -223,4 +261,4 @@ with gr.Blocks(title="Flux2 Image Generator", theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
+    demo.queue(max_size=20).launch()