Prioritize ZeroGPU init over torch import

app.py

@@ -1,5 +1,7 @@
 import os
 import random
+
+import spaces
 import gradio as gr
 import torch
 from diffusers.utils import load_image
@@ -7,7 +9,6 @@ from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipe
 from diffusers.models.controlnet_flux import FluxControlNetModel
 import numpy as np
 from huggingface_hub import login, snapshot_download
-import spaces
 
 
 # Configuration
@@ -52,9 +53,10 @@ def infer(
     seed,
     randomize_seed,
 ):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    print(f"Inference: using device: {device} (torch_dtype={torch_dtype})")
+    global DEVICE, TORCH_DTYPE
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    print(f"Inference: using device: {DEVICE} (torch_dtype={TORCH_DTYPE})")
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)