fix: ZeroGPU compatibility

core/model.py

@@ -11,7 +11,6 @@ model = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
 )
-model = model.to("cuda")
 model.eval()
 
 
@@ -38,43 +37,47 @@ def vision_infer(
     if not isinstance(images, list):
         images = [images]
 
-    pil_images = [_to_pil(img) for img in images]
+    model.to("cuda")
+    try:
+        pil_images = [_to_pil(img) for img in images]
 
-    if json_mode:
-        instruction = (
-            instruction
-            + "\n\nRespond with ONLY valid JSON. No markdown fences, no prose, no explanation."
-        )
+        if json_mode:
+            instruction = (
+                instruction
+                + "\n\nRespond with ONLY valid JSON. No markdown fences, no prose, no explanation."
+            )
 
-    content = [{"type": "image", "image": img} for img in pil_images]
-    content.append({"type": "text", "text": instruction})
+        content = [{"type": "image", "image": img} for img in pil_images]
+        content.append({"type": "text", "text": instruction})
 
-    messages = [{"role": "user", "content": content}]
+        messages = [{"role": "user", "content": content}]
 
-    inputs = processor.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        downsample_mode="16x",
-        max_slice_nums=36,
-    ).to(model.device)
-
-    with torch.no_grad():
-        generated_ids = model.generate(
-            **inputs,
+        inputs = processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
             downsample_mode="16x",
-            max_new_tokens=max_tokens,
-            do_sample=False,
-        )
+            max_slice_nums=36,
+        ).to(model.device)
+
+        with torch.no_grad():
+            generated_ids = model.generate(
+                **inputs,
+                downsample_mode="16x",
+                max_new_tokens=max_tokens,
+                do_sample=False,
+            )
 
-    trimmed = [
-        out_ids[len(in_ids):]
-        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    return processor.batch_decode(
-        trimmed,
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=False,
-    )[0]
+        trimmed = [
+            out_ids[len(in_ids):]
+            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        return processor.batch_decode(
+            trimmed,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )[0]
+    finally:
+        model.to("cpu")