Fix CUDA initialization for HF Spaces Stateless GPU environment

app.py

@@ -23,8 +23,6 @@ except ImportError:
     pass
 
 import gradio as gr
-import torch
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 # Try to import MCPClient with fallback
 try:
@@ -78,6 +76,7 @@ gpt_oss_tokenizer = None
 gpt_oss_model = None
 _initialized = False
 _init_lock = asyncio.Lock()
+_model_loading_lock = asyncio.Lock()
 
 def _current_system_prompt(style: str) -> str:
     """Get the system prompt with style suffix."""
@@ -100,29 +99,59 @@ def get_mcp_client(model_id: str, provider: str, api_key: Optional[str]) -> MCPC
         mcp_client = MCPClient(model=model_id, provider=provider, api_key=api_key)
     return mcp_client
 
-def get_gpt_oss_model_and_tokenizer():
-    """Get or create GPT-OSS-20B model and tokenizer."""
+async def get_gpt_oss_model_and_tokenizer():
+    """Get or create GPT-OSS-20B model and tokenizer with proper CUDA handling."""
     global gpt_oss_tokenizer, gpt_oss_model
-    if gpt_oss_tokenizer is None or gpt_oss_model is None:
+    
+    # Check if already loaded
+    if gpt_oss_tokenizer is not None and gpt_oss_model is not None:
+        return gpt_oss_tokenizer, gpt_oss_model
+    
+    # Use lock to prevent multiple simultaneous loads
+    async with _model_loading_lock:
+        # Double-check after acquiring lock
+        if gpt_oss_tokenizer is not None and gpt_oss_model is not None:
+            return gpt_oss_tokenizer, gpt_oss_model
+        
         try:
+            # Import here to avoid CUDA initialization in main process
+            import torch
+            from transformers import AutoTokenizer, AutoModelForCausalLM
+            
             print("🔄 Loading GPT-OSS-20B tokenizer...")
             gpt_oss_tokenizer = AutoTokenizer.from_pretrained(
                 "openai/gpt-oss-20b",
                 trust_remote_code=True,
             )
+            
             print("🔄 Loading GPT-OSS-20B model...")
+            # For HF Spaces with Stateless GPU, use specific device mapping
+            device_map = "auto"
+            if os.environ.get("SPACE_ZERO_GPU"):
+                device_map = "cpu"  # Force CPU for ZeroGPU spaces
+            
             gpt_oss_model = AutoModelForCausalLM.from_pretrained(
                 "openai/gpt-oss-20b",
-                torch_dtype="auto",
-                device_map="auto",
+                torch_dtype=torch.float16,  # Use float16 for memory efficiency
+                device_map=device_map,
                 trust_remote_code=True,
                 low_cpu_mem_usage=True,
+                # Disable gradient computation for inference
+                torch_dtype=torch.float16,
             )
+            
+            # Set model to evaluation mode
+            gpt_oss_model.eval()
+            
             print("✅ GPT-OSS-20B loaded successfully!")
+            return gpt_oss_tokenizer, gpt_oss_model
+            
         except Exception as e:
             print(f"❌ Failed to load GPT-OSS-20B: {e}")
+            # Reset globals on error
+            gpt_oss_tokenizer = None
+            gpt_oss_model = None
             raise e
-    return gpt_oss_tokenizer, gpt_oss_model
 
 async def ensure_mcp_init(model_id: str, provider: str, api_key: Optional[str]):
     """Initialize MCP server connection."""
@@ -208,7 +237,8 @@ async def stream_answer(
     # Handle GPT-OSS-20B
     if USE_GPT_OSS:
         try:
-            tokenizer, model = get_gpt_oss_model_and_tokenizer()
+            # Lazy load model only when needed
+            tokenizer, model = await get_gpt_oss_model_and_tokenizer()
             
             # Convert messages to GPT-OSS format with reasoning
             gpt_oss_messages = []
@@ -228,18 +258,23 @@ async def stream_answer(
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
-            ).to(model.device)
+            )
+            
+            # Move inputs to model device
+            if hasattr(model, 'device'):
+                inputs = {k: v.to(model.device) if hasattr(v, 'to') else v for k, v in inputs.items()}
             
             # Generate with timeout protection
             try:
-                outputs = model.generate(
-                    **inputs, 
-                    max_new_tokens=512,
-                    do_sample=True,
-                    temperature=0.7,
-                    pad_token_id=tokenizer.eos_token_id,
-                    max_time=60.0,  # 60 second timeout
-                )
+                with torch.no_grad():  # Disable gradients for inference
+                    outputs = model.generate(
+                        **inputs, 
+                        max_new_tokens=512,
+                        do_sample=True,
+                        temperature=0.7,
+                        pad_token_id=tokenizer.eos_token_id,
+                        max_time=60.0,  # 60 second timeout
+                    )
             except Exception as gen_error:
                 yield {
                     "delta": f"❌ Generation Error: {str(gen_error)}",