Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from transformers import AutoProcessor, LlavaForConditionalGeneration
 from PIL import Image
 import torch
-import numpy as np
 import spaces
 
 # Load the Llava model and processor
@@ -15,28 +14,15 @@ def generate_text(input_text="", image=None):
     if image is None:
         return "Please upload an image."
 
-    # Convert the uploaded image to a NumPy array and then to PIL format
-    image_np = np.array(image)
-    image_pil = Image.fromarray(image_np.astype('uint8'), 'RGB')
-
-    # Resize the image to the expected resolution (336 x 336)
-    image_pil = image_pil.resize((336, 336))
+    # Resize the image to the expected resolution (adjust size if necessary)
+    image = image.resize((336, 336))
 
     # Use a default prompt if no text is provided
     if not input_text:
         input_text = "Describe the image."
 
     # Prepare inputs
-    inputs = processor(text=input_text, images=image_pil, return_tensors="pt").to("cuda")
-
-    # Debug: Print the keys and types of the inputs dictionary
-    print("Processor output keys:", inputs.keys())
-    for key, value in inputs.items():
-        print(f"{key}: {type(value)}")
-
-    # Check if image tokens are generated
-    if 'input_ids' not in inputs or inputs['input_ids'].numel() == 0:
-        return "Error: No image tokens generated. Please check the image and try again."
+    inputs = processor(text=input_text, images=image, return_tensors="pt").to("cuda")
 
     # Generate output
     outputs = model.generate(**inputs)