Update app.py

app.py

@@ -127,37 +127,23 @@ def image_generator(prompt: str) -> str:
     return f"Image saved at {output_path}"
 
 # -------------------- Local LLM (Replaces HfApiModel) --------------------
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 class LocalModel:
-    """
-    Minimal local model interface compatible with smolagents CodeAgent.
-    """
     def __init__(self):
         model_name = "openlm-research/open_llama_3b"
-
-        # Load tokenizer with use_fast=False to avoid SentencePiece conversion error
         self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-
-        # Load model with appropriate dtype and device map
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             device_map="auto" if torch.cuda.is_available() else None,
         )
 
-        # Create pipeline with explicit tokenizer and model
-        self.pipeline = pipeline(
-            "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            device=0 if torch.cuda.is_available() else -1,
-        )
-
-    def generate(self, prompt, **kwargs):
-        result = self.pipeline(prompt, max_new_tokens=500, do_sample=True, **kwargs)
-        return result[0]['generated_text']
+    def generate(self, prompt, max_new_tokens=500):
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+        return self.tokenizer.decode(output[0], skip_special_tokens=True)
 
     def __call__(self, prompt, **kwargs):
         return self.generate(prompt, **kwargs)