yes

app.py

@@ -32,7 +32,7 @@ except ImportError:
         print("⚠ Triton not configured for MX - run install.sh")
 
 # ===== MAIN IMPORTS =====
-import os, gc, json, torch, warnings, traceback
+import os, gc, json, warnings, traceback
 import subprocess, sys
 from dataclasses import dataclass
 from typing import List, Dict, Optional, Any, Union
@@ -42,6 +42,9 @@ import spaces  # required for ZeroGPU
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 import numpy as np
 
+# IMPORTANT: Don't import torch at module level for ZeroGPU
+# It will be imported inside GPU-decorated functions
+
 # Suppress warnings
 warnings.filterwarnings("ignore", message=".*microscaling.*")
 warnings.filterwarnings("ignore", message=".*mx.*")
@@ -169,6 +172,8 @@ def detect_mx_format(model) -> bool:
 
 def load_base_model(device_map: Optional[str] = "auto") -> AutoModelForCausalLM:
     """Load the base model with proper MX format handling."""
+    import torch  # Import torch here for ZeroGPU compatibility
+    
     print(f"\n{'='*50}")
     print(f"Loading model: {MODEL_ID}")
     print(f"MX Format Available: {_HAS_TRITON_KERNELS}")
@@ -198,6 +203,8 @@ def load_base_model(device_map: Optional[str] = "auto") -> AutoModelForCausalLM:
         else:
             print("⚠ No triton_kernels - falling back to bf16 (dequantized)")
             print("  This will likely cause LoRA compatibility issues!")
+            # Load the model - torch imported inside function
+            import torch
             load_kwargs["torch_dtype"] = torch.bfloat16
             
             # Explicitly disable MX
@@ -205,6 +212,7 @@ def load_base_model(device_map: Optional[str] = "auto") -> AutoModelForCausalLM:
             os.environ["FORCE_MX_QUANTIZATION"] = "0"
     else:
         # Non-GPT-OSS models
+        import torch
         load_kwargs["torch_dtype"] = torch.bfloat16
     
     try:
@@ -240,6 +248,7 @@ def load_base_model(device_map: Optional[str] = "auto") -> AutoModelForCausalLM:
             
             # Try to load without MX as fallback
             print("Attempting to load model without MX format...")
+            import torch
             load_kwargs["torch_dtype"] = torch.bfloat16
             os.environ["FORCE_MX_QUANTIZATION"] = "0"
             model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **load_kwargs)
@@ -425,6 +434,8 @@ def generate_on_gpu(
     seed: Optional[int]
 ) -> Dict[str, str]:
     """Run generation on GPU."""
+    import torch  # Import torch inside GPU function for ZeroGPU
+    
     try:
         # Set seed if provided
         if seed is not None:
@@ -441,6 +452,7 @@ def generate_on_gpu(
         model.eval()
         
         # Prepare inputs
+        import torch  # Make sure torch is available
         device = next(model.parameters()).device
         
         if HARMONY_AVAILABLE and isinstance(prompt, list):
@@ -492,6 +504,7 @@ def generate_on_gpu(
     
     finally:
         # Cleanup
+        import torch
         if 'model' in locals():
             del model
         gc.collect()