Update model_manager.py

model_manager.py

@@ -58,6 +58,7 @@ class LazyLlamaModel:
     
     def __init__(self):
         """Initialize only once"""
+        # Added hasattr check to prevent re-initialization
         if hasattr(self, '_initialized') and self._initialized:
             return
         
@@ -96,6 +97,11 @@ class LazyLlamaModel:
                 trust_remote_code=True,
             )
             
+            # Add validation after tokenizer load
+            if self.tokenizer is None:
+                raise RuntimeError("Tokenizer failed to load")
+            logger.info(f"✓ Tokenizer loaded: {type(self.tokenizer).__name__}")
+            
             # Load model
             self.model = AutoModelForCausalLM.from_pretrained(
                 LLAMA_MODEL_ID,
@@ -106,8 +112,14 @@ class LazyLlamaModel:
                 torch_dtype=torch.bfloat16,
             )
             
+            # Add validation after model load
+            if self.model is None:
+                raise RuntimeError("Model failed to load")
+            logger.info(f"✓ Model loaded: {type(self.model).__name__}")
+            
             # Create pipeline
-            self.pipe = pipeline(
+            # Store pipeline reference explicitly
+            pipeline_obj = pipeline(
                 "text-generation",
                 model=self.model,
                 tokenizer=self.tokenizer,
@@ -115,15 +127,37 @@ class LazyLlamaModel:
                 device_map="auto",
             )
             
+            # FIXED: Validate pipeline before assignment
+            if pipeline_obj is None:
+                raise RuntimeError("Pipeline creation returned None")
+            
+            if not callable(pipeline_obj):
+                raise RuntimeError(f"Pipeline is not callable: {type(pipeline_obj)}")
+            
+            # Assign to instance
+            self.pipe = pipeline_obj
+            
+            # FIXED: Double-check assignment succeeded
+            if self.pipe is None:
+                raise RuntimeError("Pipeline assignment failed - pipe is still None")
+            
+            logger.info(f"✓ Pipeline created and verified: {type(self.pipe).__name__}")
+            
             logger.info("="*60)
             logger.info("✅ MODEL LOADED & CACHED")
             logger.info(f"  Model: {LLAMA_MODEL_ID}")
+            logger.info(f"  Tokenizer: {type(self.tokenizer).__name__}")
+            logger.info(f"  Pipeline: {type(self.pipe).__name__}")
             logger.info(f"  Memory: ~1GB VRAM")
             logger.info(f"  Context: 128K tokens")
             logger.info("="*60)
             
         except Exception as e:
             logger.error(f"Failed to load model: {e}")
+            # Clean up partial state
+            self.model = None
+            self.tokenizer = None
+            self.pipe = None
             raise
     
     def generate(
@@ -151,12 +185,34 @@ class LazyLlamaModel:
         if self.model is None:
             self._load_model()
         
+        # Add comprehensive validation after loading
+        if self.pipe is None:
+            # Try reloading if pipe is None but model exists
+            if self.model is not None:
+                logger.warning("Pipeline is None but model exists - attempting reload")
+                self.model = None  # Force full reload
+                self._load_model()
+            
+            # If still None, fail with clear error
+            if self.pipe is None:
+                raise RuntimeError(
+                    "Pipeline is None after loading. This may be a ZeroGPU context issue. "
+                    "Check that _load_model() completed successfully."
+                )
+        
+        # Verify pipeline is callable
+        if not callable(self.pipe):
+            raise RuntimeError(f"Pipeline exists but is not callable: {type(self.pipe)}")
+        
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": user_message},
         ]
         
         try:
+            # FIXED: Add logging for debugging
+            logger.debug(f"Calling pipeline with {len(messages)} messages, max_tokens={max_tokens}")
+            
             outputs = self.pipe(
                 messages,
                 max_new_tokens=max_tokens,
@@ -167,11 +223,39 @@ class LazyLlamaModel:
                 repetition_penalty=1.15,
             )
             
-            result = outputs[0]["generated_text"][-1]["content"]
+            # FIXED: Validate output structure
+            if not outputs or len(outputs) == 0:
+                logger.error("Pipeline returned empty output")
+                return ""
+            
+            if not isinstance(outputs[0], dict):
+                logger.error(f"Unexpected output format: {type(outputs[0])}")
+                return ""
+            
+            if "generated_text" not in outputs[0]:
+                logger.error(f"No 'generated_text' in output: {outputs[0].keys()}")
+                return ""
+            
+            generated = outputs[0]["generated_text"]
+            
+            if not generated or len(generated) == 0:
+                logger.error("Generated text is empty")
+                return ""
+            
+            # Extract final message content
+            result = generated[-1]["content"]
+            
+            logger.debug(f"Generated {len(result)} characters successfully")
+            
             return result.strip()
             
         except Exception as e:
             logger.error(f"Generation error: {e}")
+            logger.error(f"Error type: {type(e).__name__}")
+            logger.error(f"Pipeline type: {type(self.pipe)}")
+            logger.error(f"Pipeline callable: {callable(self.pipe)}")
+            import traceback
+            logger.error(traceback.format_exc())
             return ""
     
     def generate_streaming(
@@ -191,6 +275,17 @@ class LazyLlamaModel:
         if self.model is None:
             self._load_model()
         
+        # FIXED: Add validation for streaming
+        if self.model is None:
+            logger.error("Model is None in generate_streaming")
+            yield ""
+            return
+        
+        if self.tokenizer is None:
+            logger.error("Tokenizer is None in generate_streaming")
+            yield ""
+            return
+        
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": user_message},
@@ -231,24 +326,29 @@ class LazyLlamaModel:
                 
         except Exception as e:
             logger.error(f"Streaming error: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
             yield ""
     
     def is_loaded(self) -> bool:
         """Check if model is loaded"""
-        return self.model is not None
+        return self.model is not None and self.pipe is not None
     
     def get_model_info(self) -> dict:
         """Get model information"""
         return {
             "model_id": LLAMA_MODEL_ID,
             "loaded": self.is_loaded(),
+            "model_exists": self.model is not None,
+            "tokenizer_exists": self.tokenizer is not None,
+            "pipe_exists": self.pipe is not None,
+            "pipe_callable": callable(self.pipe) if self.pipe else False,
             "quantization": "4-bit NF4",
             "size_gb": 1.0,
             "context_length": 128000,
             "lazy_loading": True,
         }
 
-
 # Global instance - model loads on first use
 _model_instance = None