updated generate text for image

app.py

@@ -70,43 +70,59 @@ def prepare_inputs(text_list, image_input, image_processor_instance, tokenizer_i
     attention_mask = processed_text.attention_mask.to(device_to_use)
     return {"pixel_values": processed_image, "input_ids": input_ids, "attention_mask": attention_mask}
 
-def generate_text_for_image(image_input, prompt_input):
+from typing import Optional
+from PIL import Image as PILImage # Add at the top of your app.py
+
+# ... (other imports and model loading) ...
+
+def generate_text_for_image(image_input: Optional[PILImage.Image], prompt_input: Optional[str]) -> str:
     if model is None or image_processor is None or tokenizer is None:
         return "Error: Model or processor components not loaded correctly. Check logs."
     if image_input is None: return "Please upload an image."
     if not prompt_input: return "Please provide a prompt."
 
     try:
-        if not isinstance(image_input, Image.Image):
-            pil_image = Image.fromarray(image_input)
-        else:
-            pil_image = image_input
-        if pil_image.mode != "RGB": pil_image = pil_image.convert("RGB")
+        current_pil_image = image_input
+        if not isinstance(current_pil_image, PILImage.Image):
+             current_pil_image = PILImage.fromarray(current_pil_image)
+        if current_pil_image.mode != "RGB":
+            current_pil_image = current_pil_image.convert("RGB")
 
         inputs = prepare_inputs(
-            text_list=[prompt_input], image_input=pil_image,
+            text_list=[prompt_input], image_input=current_pil_image,
             image_processor_instance=image_processor, tokenizer_instance=tokenizer, device_to_use=device
         )
         
+        print(f"Debug: Passing to model.generate: pixel_values_shape={inputs['pixel_values'].shape}, input_ids_shape={inputs['input_ids'].shape}, attention_mask_shape={inputs['attention_mask'].shape}")
+
+        # Call model.generate with positional arguments matching nanoVLM's VisionLanguageModel.generate
         generated_ids = model.generate(
-            pixel_values=inputs['pixel_values'], input_ids=inputs['input_ids'],
-            attention_mask=inputs['attention_mask'], max_new_tokens=150, num_beams=3,
-            no_repeat_ngram_size=2, early_stopping=True, pad_token_id=tokenizer.pad_token_id
+            inputs['pixel_values'],       # pixel_values
+            inputs['input_ids'],          # prompt_token_ids
+            inputs['attention_mask'],     # attention_mask
+            150                           # max_new_tokens (as a positional argument)
+            # You can add temperature=..., top_k=... here if desired, as they are keyword args in nanoVLM's generate
         )
         
         generated_text_list = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
         generated_text = generated_text_list[0] if generated_text_list else ""
 
+        # Clean up prompt if it's echoed (optional, depends on model behavior)
         if prompt_input and generated_text.startswith(prompt_input):
              cleaned_text = generated_text[len(prompt_input):].lstrip(" ,.:")
         else:
             cleaned_text = generated_text
+
         return cleaned_text.strip()
+
     except Exception as e:
         print(f"Error during generation: {e}")
-        import traceback; traceback.print_exc()
+        import traceback
+        traceback.print_exc()
         return f"An error occurred during text generation: {str(e)}"
 
+# ... (rest of app.py)
+
 description = "Interactive demo for lusxvr/nanoVLM-222M."
 # example_image_url = "http://images.cocodataset.org/val2017/000000039769.jpg" # Not used for now