loading chain

app.py

@@ -1,3 +1,6 @@
+# Import spaces first to avoid CUDA initialization conflicts
+import spaces
+
 import gradio as gr
 import numpy as np
 import random
@@ -5,15 +8,16 @@ import torch
 from PIL import Image
 from torchvision import transforms
 from diffusers import DiffusionPipeline, AutoencoderKL
-import spaces
 
 # Define constants
 flux_dtype = torch.bfloat16
 vae_dtype = torch.float32
-device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+# Move device selection after spaces import
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 def load_models():
     # Load the initial VAE model for preprocessing in float32
     vae_model_name = "runwayml/stable-diffusion-v1-5"
@@ -28,7 +32,13 @@ def load_models():
     
     return vae, pipe
 
-vae, pipe = load_models()
+# Defer model loading until it's needed
+vae, pipe = None, None
+
+def ensure_models_loaded():
+    global vae, pipe
+    if vae is None or pipe is None:
+        vae, pipe = load_models()
 
 def preprocess_image(image, image_size):
     preprocess = transforms.Compose([
@@ -52,6 +62,8 @@ def encode_image(image, vae):
 
 @spaces.GPU()
 def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    ensure_models_loaded()
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
@@ -113,6 +125,8 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         print(f"Error during inference: {e}")
         return fallback_image, seed
 
+# ... (rest of the Gradio interface code remains the same)
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",