y

app.py

@@ -4,7 +4,6 @@ import gradio as gr
 import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
-import spaces
 from transformers import T5Tokenizer, T5EncoderModel
 from diffusers import StableDiffusionXLPipeline, DDIMScheduler, EulerDiscreteScheduler, DPMSolverMultistepScheduler
 from safetensors.torch import load_file
@@ -13,20 +12,14 @@ from two_stream_shunt_adapter import TwoStreamShuntAdapter
 from configs import T5_SHUNT_REPOS
 
 # ─── Device & Model Setup ─────────────────────────────────────
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-# T5 Model for semantic understanding
-t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base")
-t5_mod = T5EncoderModel.from_pretrained("google/flan-t5-base").to(device).eval()
+# Don't initialize CUDA here for ZeroGPU compatibility
+device = None  # Will be set inside the GPU function
+dtype = torch.float16
 
-# SDXL Pipeline with proper text encoders
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=dtype,
-    variant="fp16" if dtype == torch.float16 else None,
-    use_safetensors=True
-).to(device)
+# Don't load models here - will load inside GPU function
+t5_tok = None
+t5_mod = None 
+pipe = None
 
 # Available schedulers
 SCHEDULERS = {
@@ -47,6 +40,7 @@ config_g = T5_SHUNT_REPOS["clip_g"]["config"]
 from safetensors.torch import safe_open
 
 def load_adapter(repo, filename, config):
+    # Don't initialize device here
     path = hf_hub_download(repo_id=repo, filename=filename)
     
     model = TwoStreamShuntAdapter(config).eval()
@@ -55,7 +49,7 @@ def load_adapter(repo, filename, config):
         for key in f.keys():
             tensors[key] = f.get_tensor(key)
     model.load_state_dict(tensors)
-    model.to(device)
+    # Device will be set when called from GPU function
     return model
 
 # ─── Visualization ────────────────────────────────────────────
@@ -135,11 +129,34 @@ def encode_sdxl_prompt(prompt, negative_prompt=""):
     }
 
 # ─── Inference ────────────────────────────────────────────────
+@torch.no_grad()
+def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noise, gate_prob, 
+          use_anchor, steps, cfg_scale, scheduler_name, width, height, seed):
+    
+# ─── Inference ────────────────────────────────────────────
 @spaces.GPU
 @torch.no_grad()
 def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noise, gate_prob, 
           use_anchor, steps, cfg_scale, scheduler_name, width, height, seed):
     
+    # Initialize device and models inside GPU context
+    global t5_tok, t5_mod, pipe
+    device = torch.device("cuda")
+    dtype = torch.float16
+    
+    # Load models if not already loaded
+    if t5_tok is None:
+        t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base")
+        t5_mod = T5EncoderModel.from_pretrained("google/flan-t5-base").to(device).eval()
+        
+    if pipe is None:
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            torch_dtype=dtype,
+            variant="fp16",
+            use_safetensors=True
+        ).to(device)
+    
     # Set seed for reproducibility
     if seed != -1:
         torch.manual_seed(seed)
@@ -168,8 +185,8 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
     print(f"CLIP-G shape: {clip_embeds['clip_g'].shape}")
     
     # Load adapters
-    adapter_l = load_adapter(repo_l, adapter_l_file, config_l) if adapter_l_file else None
-    adapter_g = load_adapter(repo_g, adapter_g_file, config_g) if adapter_g_file else None
+    adapter_l = load_adapter(repo_l, adapter_l_file, config_l).to(device) if adapter_l_file else None
+    adapter_g = load_adapter(repo_g, adapter_g_file, config_g).to(device) if adapter_g_file else None
     
     # Apply CLIP-L adapter
     if adapter_l is not None: