Update app.py

app.py

@@ -8,59 +8,70 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Load model & tokenizer
+# Global variables for model and tokenizer (lazy loading)
+model = None
+tokenizer = None
 MODEL_NAME = "ubiodee/Test_Plutus"
+FALLBACK_TOKENIZER = "gpt2"
 
+# Load tokenizer at startup (lightweight, no model yet)
 try:
-    logger.info("Loading tokenizer with use_fast=False...")
+    logger.info("Loading tokenizer at startup with legacy versions...")
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        use_fast=False,  # Use slow tokenizer to avoid fast tokenizer errors
-        use_safetensors=True,
-        trust_remote_code=True,  # Allow custom tokenizer code
+        use_fast=False,
+        trust_remote_code=True,
     )
-    logger.info("Tokenizer loaded successfully.")
+    logger.info("Primary tokenizer loaded successfully.")
 except Exception as e:
-    logger.error(f"Tokenizer loading failed: {str(e)}")
-    raise
-
-try:
-    logger.info("Loading model with 8-bit quantization...")
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map="auto",  # Automatically map to GPU/CPU
-        load_in_8bit=True,  # Use 8-bit quantization to match model
-        torch_dtype=torch.bfloat16,  # Use bfloat16 for efficiency
-        use_safetensors=True,
-        low_cpu_mem_usage=True,  # Reduce CPU memory during loading
-        trust_remote_code=True,  # Allow custom model code
+    logger.warning(f"Primary tokenizer failed: {str(e)}. Using fallback.")
+    tokenizer = AutoTokenizer.from_pretrained(
+        FALLBACK_TOKENIZER,
+        use_fast=False,
+        trust_remote_code=True,
     )
-    model.eval()
-    logger.info("Model loaded successfully.")
-except Exception as e:
-    logger.error(f"Model loading failed: {str(e)}")
-    raise
+    logger.info("Fallback tokenizer loaded.")
 
-# Set pad token if not defined
+# Set pad token
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
     logger.info("Set pad_token_id to eos_token_id.")
 
-# Move model to GPU if available
-if torch.cuda.is_available():
-    model.to("cuda")
-    logger.info("Model moved to GPU.")
-else:
-    logger.warning("No GPU available, using CPU.")
+def load_model():
+    """Load model inside GPU context to enable quantization."""
+    global model
+    if model is None:
+        try:
+            logger.info("Loading model with CPU fallback (full precision)...")
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                torch_dtype=torch.float16,  # Use fp16 for memory efficiency without bitsandbytes
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+            )
+            model.eval()
+            if torch.cuda.is_available():
+                model.to("cuda")
+                logger.info("Model loaded and moved to GPU.")
+            else:
+                logger.warning("GPU not available; using CPU.")
+        except Exception as e:
+            logger.error(f"Model loading failed: {str(e)}")
+            raise
+    return model
 
-# Response function with GPU decorator
-@spaces.GPU
+# Response function: Load model on first call, then reuse
+@spaces.GPU(duration=300)  # Allow up to 5min for loading + inference
 def generate_response(prompt, progress=gr.Progress()):
-    progress(0.1, desc="Tokenizing input...")
+    global model
+    progress(0.1, desc="Loading model if needed...")
+    model = load_model()  # Ensures model is loaded in GPU context
+    
+    progress(0.3, desc="Tokenizing input...")
     try:
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         
-        progress(0.5, desc="Generating response...")
+        progress(0.6, desc="Generating response...")
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
@@ -73,7 +84,7 @@ def generate_response(prompt, progress=gr.Progress()):
             )
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Remove the prompt from the output
+        # Remove prompt from output
         if response.startswith(prompt):
             response = response[len(prompt):].strip()
         
@@ -93,4 +104,4 @@ demo = gr.Interface(
 )
 
 # Launch with queueing
-demo.queue(max_size=10).launch(enable_queue=True, max_threads=1)
+demo.queue(max_size=5).launch(enable_queue=True, max_threads=1)