Fix PyTorch CVE-2025-32434: upgrade to v2.6+, use safetensors, restore Llama 3.1

Dockerfile

@@ -22,8 +22,8 @@ WORKDIR /app
 # Copy requirements
 COPY requirements.txt .
 
-# Install PyTorch with CUDA support
-RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+# Install PyTorch 2.6+ with CUDA support to fix CVE-2025-32434
+RUN pip install torch>=2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
 
 # Skip llama-cpp-python to avoid compilation - using transformers instead
 

app.py

@@ -62,9 +62,8 @@ async def load_model():
             
             logger.info("Loading model with transformers...")
             
-            # Use a working model while waiting for Llama 3.1 access
-            # TODO: Change back to "meta-llama/Llama-3.1-8B-Instruct" once you have access
-            base_model_name = "microsoft/DialoGPT-medium"
+            # Use Llama 3.1 8B Instruct (user now has access)
+            base_model_name = "meta-llama/Llama-3.1-8B-Instruct"
             
             # Get HF token from environment
             hf_token = os.getenv("HF_TOKEN")
@@ -82,6 +81,7 @@ async def load_model():
                 device_map="auto" if device == "cuda" else None,
                 trust_remote_code=True,
                 low_cpu_mem_usage=True,
+                use_safetensors=True,  # Force safetensors to avoid CVE-2025-32434
                 token=hf_token
             )
             
@@ -92,27 +92,7 @@ async def load_model():
             
         except Exception as e:
             logger.error(f"Error loading model with transformers: {str(e)}")
-            # Fallback to a smaller, more reliable model
-            logger.info("Falling back to smaller model...")
-            
-            try:
-                base_model_name = "microsoft/DialoGPT-medium"
-                
-                tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-                model = AutoModelForCausalLM.from_pretrained(
-                    base_model_name,
-                    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-                    device_map="auto" if device == "cuda" else None
-                )
-                
-                if device == "cuda":
-                    model = model.to(device)
-                    
-                logger.info("Fallback model loaded successfully!")
-                
-            except Exception as fallback_error:
-                logger.error(f"Fallback model also failed: {str(fallback_error)}")
-                raise
+            raise # Re-raise the error to stop startup if primary model fails
             
     except Exception as e:
         logger.error(f"Error loading model: {str(e)}")

requirements.txt

@@ -1,7 +1,7 @@
 fastapi==0.104.1
 uvicorn[standard]==0.24.0
 pydantic==2.5.0
-torch>=2.0.0
+torch>=2.6.0
 transformers>=4.35.0
 accelerate>=0.24.0
 bitsandbytes>=0.41.0