Update app.py

app.py

@@ -1,19 +1,24 @@
 import spaces
-import argparse
 import os
 import time
 from os import path
-from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download
+import gradio as gr
+import torch
 
+# Set cache paths before importing transformers/diffusers
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-os.environ["TRANSFORMERS_CACHE"] = cache_path
-os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 
-import gradio as gr
-import torch
-from diffusers import FluxPipeline
+# Import with proper error handling
+try:
+    from diffusers import DiffusionPipeline
+    from diffusers.models import FluxTransformer2DModel
+    from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+    from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+except ImportError as e:
+    print(f"Import error: {e}")
+    # Fallback to DiffusionPipeline if FluxPipeline is not available
+    from diffusers import DiffusionPipeline
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
@@ -30,9 +35,34 @@ class timer:
 if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
-pipe.fuse_lora(lora_scale=0.125)
+# Use DiffusionPipeline as a more stable alternative
+try:
+    # Try to load as FluxPipeline first
+    pipe = DiffusionPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev", 
+        torch_dtype=torch.bfloat16,
+        use_safetensors=True
+    )
+except Exception as e:
+    print(f"Error loading FLUX pipeline: {e}")
+    # Fallback to a simpler configuration
+    pipe = DiffusionPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        torch_dtype=torch.float16,
+        safety_checker=None,
+        requires_safety_checker=False
+    )
+
+# Try to load LoRA weights with error handling
+try:
+    from huggingface_hub import hf_hub_download
+    lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
+    pipe.load_lora_weights(lora_path)
+    pipe.fuse_lora(lora_scale=0.125)
+except Exception as e:
+    print(f"Warning: Could not load LoRA weights: {e}")
+    print("Continuing without LoRA acceleration...")
+
 pipe.to(device="cuda", dtype=torch.bfloat16)
 
 # Define example prompts
@@ -192,15 +222,29 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     def process_image(height, width, steps, scales, prompt, seed):
         global pipe
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-            return pipe(
-                prompt=[prompt],
-                generator=torch.Generator().manual_seed(int(seed)),
-                num_inference_steps=int(steps),
-                guidance_scale=float(scales),
-                height=int(height),
-                width=int(width),
-                max_sequence_length=256
-            ).images[0]
+            try:
+                # Try the standard call
+                result = pipe(
+                    prompt=[prompt],
+                    generator=torch.Generator().manual_seed(int(seed)),
+                    num_inference_steps=int(steps),
+                    guidance_scale=float(scales),
+                    height=int(height),
+                    width=int(width),
+                    max_sequence_length=256
+                )
+            except TypeError:
+                # Fallback for different pipeline signatures
+                result = pipe(
+                    prompt=prompt,
+                    generator=torch.Generator().manual_seed(int(seed)),
+                    num_inference_steps=int(steps),
+                    guidance_scale=float(scales),
+                    height=int(height),
+                    width=int(width)
+                )
+            
+            return result.images[0]
 
     generate_btn.click(
         process_image,