Update app

app.py

@@ -17,6 +17,9 @@ import re
 MODEL_REPO = "Nishan30/n8n-workflow-generator"  # Update with your HF repo
 BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 
+# Memory optimization: Set to True for 8-bit quantization (uses less memory but slower)
+USE_8BIT = False  # Change to True if you get out-of-memory errors
+
 # ==============================================================================
 # MODEL LOADING
 # ==============================================================================
@@ -25,17 +28,35 @@ def load_model():
     """Load model once and cache it"""
     print("Loading model...")
 
-    # Load base model
+    # Prepare model loading kwargs
+    model_kwargs = {
+        "device_map": "auto",
+        "trust_remote_code": True,
+        "low_cpu_mem_usage": True,
+        "offload_folder": "offload",
+        "offload_state_dict": True,
+    }
+
+    # Use 8-bit quantization if enabled (saves memory)
+    if USE_8BIT:
+        print("Using 8-bit quantization for memory efficiency...")
+        model_kwargs["load_in_8bit"] = True
+    else:
+        model_kwargs["torch_dtype"] = torch.float16
+
+    # Load base model with memory optimization
     base_model = AutoModelForCausalLM.from_pretrained(
         BASE_MODEL,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True
+        **model_kwargs
     )
 
     # Load LoRA adapter with error handling for unsupported parameters
     try:
-        model = PeftModel.from_pretrained(base_model, MODEL_REPO)
+        model = PeftModel.from_pretrained(
+            base_model,
+            MODEL_REPO,
+            offload_folder="offload",  # Enable disk offloading for adapter too
+        )
     except TypeError as e:
         if "unexpected keyword argument" in str(e):
             print(f"⚠️ Warning: {e}")
@@ -73,7 +94,11 @@ def load_model():
                     continue
 
             # Load from temp directory
-            model = PeftModel.from_pretrained(base_model, temp_dir)
+            model = PeftModel.from_pretrained(
+                base_model,
+                temp_dir,
+                offload_folder="offload"
+            )
 
             # Cleanup
             shutil.rmtree(temp_dir)
@@ -82,6 +107,10 @@ def load_model():
 
     tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
 
+    # Set pad token if not present
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
     print("Model loaded successfully!")
     return model, tokenizer
 
@@ -488,4 +517,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         share=False
-    )
+    )