Update handler.py

handler.py

@@ -35,7 +35,8 @@ class EndpointHandler:
         image_url = inputs.get("url")
         image_data = inputs.get("image_data")
         prompt = inputs.get("prompt")
-
+        max_tokens = inputs.get("max_tokens", 100)
+        
         if not prompt:
             logs.append("S'ha de proporcionar 'prompt' en 'inputs'.")
             return {"error": "The 'prompt' must be provided in 'inputs'.", "logs": logs}
@@ -70,7 +71,7 @@ class EndpointHandler:
         try:
             logs.append("Processant imatge amb el model.")
             inputs = self.processor(prompt, image, return_tensors="pt").to("cuda")
-            output = self.model.generate(**inputs, max_new_tokens=100)
+            output = self.model.generate(**inputs, max_new_tokens=max_tokens)
             result = self.processor.decode(output[0], skip_special_tokens=True)
             logs.append("Processament complet.")
             return {"input_prompt": prompt, "model_output": result, "logs": logs}