ZeroGPU: move models to CUDA inside @gpu calls, not in the main process

ZeroGPU forbids real CUDA init outside an @gpu allocation. load() now keeps
models on CPU (no .cuda(), no device_map=auto, no pipe.to('cuda')); the move to
CUDA happens inside _minicpm_generate/_nemotron_generate/_run_art_pipe.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

art.py

@@ -65,9 +65,9 @@ class DiffusersImageClient:
         import torch
         from diffusers import AutoPipelineForText2Image
 
+        # Stay on CPU here: on ZeroGPU the move to CUDA happens in _run_art_pipe.
         pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=best_torch_dtype(torch))
         pipe.set_progress_bar_config(disable=True)
-        pipe = move_pipe_to_device(pipe, torch)
         return cls(pipe, steps=steps, guidance_scale=guidance_scale, width=width, height=height)
 
 
@@ -78,6 +78,7 @@ _art_pipe: Any = None
 # so the forked GPU worker inherits it; returns a data URI string (picklable).
 @gpu
 def _run_art_pipe(prompt: str, steps: int, guidance_scale: float, width: int, height: int) -> str:
+    _art_pipe.to("cuda")
     result = _art_pipe(
         prompt=prompt,
         num_inference_steps=steps,

local_llm.py

@@ -99,8 +99,10 @@ class NemotronTransformersChatClient:
     ) -> "NemotronTransformersChatClient":  # pragma: no cover
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
+        # Load on CPU (no device_map="auto"): on ZeroGPU the move to CUDA must
+        # happen inside the @gpu call, in _nemotron_generate.
         tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForCausalLM.from_pretrained(model_path, **transformers_model_kwargs())
+        model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto")
         return cls(model.eval(), tokenizer, max_new_tokens=max_new_tokens, temperature=temperature)
 
 
@@ -112,6 +114,7 @@ _nemotron_tokenizer: Any = None
 # globals so the forked GPU worker inherits it (only strings cross the boundary).
 @gpu
 def _nemotron_generate(system: str, user: str, max_new_tokens: int, temperature: float) -> str:
+    _nemotron_model.to("cuda")
     messages = chat_messages(system, user)
     inputs = _nemotron_tokenizer.apply_chat_template(
         messages,
@@ -181,8 +184,8 @@ class MiniCPMTransformersChatClient:
             attn_implementation="sdpa",
             torch_dtype=local_torch_dtype(torch),
         )
-        if torch.cuda.is_available():
-            model = model.cuda()
+        # Stay on CPU here: on ZeroGPU the GPU only exists inside @gpu calls,
+        # so the move to CUDA happens in _minicpm_generate.
         tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
         return cls(model.eval(), tokenizer, max_new_tokens=max_new_tokens, temperature=temperature)
 
@@ -195,6 +198,7 @@ _minicpm_tokenizer: Any = None
 # globals so the forked GPU worker inherits it (only strings cross the boundary).
 @gpu
 def _minicpm_generate(system: str, user: str, max_new_tokens: int, temperature: float) -> str:
+    _minicpm_model.to("cuda")
     return str(
         _minicpm_model.chat(
             msgs=[{"role": "user", "content": user}],

tests/test_local_llm.py

@@ -72,6 +72,10 @@ class FakeTokenizer:
 class FakeModel:
     device = "cuda"
 
+    # No-op device move (ZeroGPU does the real one inside the @gpu call).
+    def to(self, device: str) -> "FakeModel":
+        return self
+
     # Return a generated token sequence with prompt prefix included.
     def generate(self, inputs: FakeInputs, **kwargs: Any) -> list[list[int]]:
         self.inputs = inputs
@@ -80,6 +84,10 @@ class FakeModel:
 
 
 class FakeMiniCPMModel:
+    # No-op device move (ZeroGPU does the real one inside the @gpu call).
+    def to(self, device: str) -> "FakeMiniCPMModel":
+        return self
+
     # Capture MiniCPM chat kwargs.
     def chat(self, **kwargs: Any) -> str:
         self.kwargs = kwargs