Update app.py

app.py

@@ -15,7 +15,7 @@ except ImportError:
 import soundfile as sf
 import traceback
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 from matcha.models.matcha_tts import MatchaTTS
 from matcha.hifigan.models import Generator as HiFiGAN
@@ -55,8 +55,17 @@ def load_models():
 TRANSLATION_BASE_MODEL = "sarvamai/sarvam-translate"
 TRANSLATION_ADAPTER = "GAASH-Lab/Sarvam-Kashmiri-finetuned"
 
+# Global cache for translation model (loaded lazily when GPU is available)
+_trans_cache = {"tokenizer": None, "model": None, "loaded": False}
+
 def load_translation_models():
-    print("[*] Loading Sarvam Translate Adapter...")
+    """Load translation model lazily on first use (CPU deployment)."""
+    global _trans_cache
+    
+    if _trans_cache["loaded"]:
+        return _trans_cache["tokenizer"], _trans_cache["model"]
+    
+    print("[*] Loading Sarvam Translate Adapter (CPU mode)...")
     try:
         # Load the tokenizer with left padding (required for causal LM)
         tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_BASE_MODEL, trust_remote_code=True)
@@ -64,21 +73,14 @@ def load_translation_models():
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
-        # Use 4-bit quantization to fit in 16GB memory limit
-        print("[*] Using 4-bit quantization to reduce memory usage...")
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4"
-        )
-        
-        # Load the base model with 4-bit quantization
-        print("[*] Loading base model as AutoModelForCausalLM (4-bit)...")
+        # Load the base model on CPU with bfloat16 to reduce memory
+        # bfloat16 is better supported on CPU than float16
+        print("[*] Loading base model on CPU (bfloat16)...")
         base_model = AutoModelForCausalLM.from_pretrained(
-            TRANSLATION_BASE_MODEL, 
-            quantization_config=quantization_config,
-            device_map="auto",
+            TRANSLATION_BASE_MODEL,
+            torch_dtype=torch.bfloat16,
+            device_map="cpu",
+            low_cpu_mem_usage=True,
             trust_remote_code=True
         )
         
@@ -87,18 +89,25 @@ def load_translation_models():
         model = PeftModel.from_pretrained(base_model, TRANSLATION_ADAPTER)
         model.eval()
         
-        print(f"[+] Translation model loaded successfully.")
+        print(f"[+] Translation model loaded successfully on CPU.")
+        _trans_cache["tokenizer"] = tokenizer
+        _trans_cache["model"] = model
+        _trans_cache["loaded"] = True
         return tokenizer, model
     except Exception as e:
         print(f"[-] Error loading translation model: {e}")
         traceback.print_exc()
         return None, None
 
+# Load TTS models at startup (they're smaller)
 model, vocoder = load_models()
-trans_tokenizer, trans_model = load_translation_models()
+# Translation model will be loaded lazily when GPU is available
 
 def _translate_impl(text):
     """Internal translation implementation - matching evaluate_model.py approach."""
+    # Load model lazily (will be cached after first load)
+    trans_tokenizer, trans_model = load_translation_models()
+    
     if trans_model is None:
         return "Translation model unavailable."
     
@@ -113,7 +122,7 @@ def _translate_impl(text):
         prompt = trans_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         inputs = trans_tokenizer(prompt, padding=True, truncation=True, max_length=512, return_tensors="pt")
         
-        # Move inputs to model's device (handles device_map="auto")
+        # Move inputs to model's device
         inputs = {k: v.to(trans_model.device) for k, v in inputs.items()}
         
         print(f"[DEBUG] Input tokens: {inputs['input_ids'].shape[1]}")
@@ -155,14 +164,9 @@ def _translate_impl(text):
         traceback.print_exc()
         return "Error during translation generation."
 
-# Wrap with GPU decorator if available
-if SPACES_AVAILABLE:
-    @spaces.GPU
-    def translate(text):
-        return _translate_impl(text)
-else:
-    def translate(text):
-        return _translate_impl(text)
+# Simple wrapper function for CPU deployment
+def translate(text):
+    return _translate_impl(text)
 
 
 # --- Update the function signature to accept two arguments ---