Update app.py

app.py

@@ -2,6 +2,9 @@ import os
 # Set environment variable before importing torch to avoid nested tensor issues
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 
+# Import spaces FIRST before any torch imports
+import spaces
+
 import time
 import gradio as gr
 import torch
@@ -22,9 +25,6 @@ from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
 
-# Import spaces after other imports to minimize conflicts
-import spaces
-
 # ---------------- Encoders ----------------
 
 class HFEmbedder(nn.Module):
@@ -62,7 +62,6 @@ class HFEmbedder(nn.Module):
         return outputs[self.output_key]
 
 # Initialize models without GPU decorator first
-device = "cuda" if torch.cuda.is_available() else "cpu"
 t5 = None
 clip = None
 ae = None
@@ -73,6 +72,7 @@ def initialize_models():
     global t5, clip, ae, model, model_initialized
     if not model_initialized:
         print("Initializing models...")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
         clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
         ae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)