Update app.py

app.py

@@ -1,8 +1,25 @@
 import os
 import re
 import gradio as gr
+import spaces
+
+# CRITICAL: Disable PyTorch compiler BEFORE importing torch
+os.environ["PYTORCH_NO_CUDA_MEMORY_CACHING"] = "1"
+os.environ["TORCH_COMPILE_DISABLE"] = "1"
+os.environ["TORCH_INDUCTOR_DISABLE"] = "1"
+os.environ["TORCHINDUCTOR_DISABLE_CUDAGRAPHS"] = "1"
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+os.environ["TORCH_USE_CUDA_DSA"] = "0"
+
+# Now import torch and disable its compiler features
 import torch
-import spaces  # Import spaces module
+if hasattr(torch, "_dynamo"):
+    if hasattr(torch._dynamo, "config"):
+        torch._dynamo.config.suppress_errors = True
+    if hasattr(torch._dynamo, "disable"):
+        torch._dynamo.disable()
+        print("Disabled torch._dynamo")
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Global variables for model and tokenizer
@@ -202,8 +219,9 @@ def load_models():
             token=hf_token,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             device_map=device_map,
-            # Ensure we use the model's full capabilities
-            use_cache=True
+            use_cache=True,
+            use_flash_attention_2=False,
+            _attn_implementation="eager"  # Use eager mode to avoid compiler issues
         )
         
         # Cache the loaded model and tokenizer