Uploading FoodExtract-Vision demo app.py from YouTube tutorial video

app.py

@@ -1,53 +1,65 @@
 import torch
 import gradio as gr
-
 import spaces
 from transformers import pipeline
 
 BASE_MODEL_ID = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
-# FINE_TUNED_MODEL_ID = "mrdbourke/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1"
 FINE_TUNED_MODEL_ID = "ninjals/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1-VIDEO"
 OUTPUT_TOKENS = 256
-
-# Load original base model (no fine-tuning)
-print(f"[INFO] Loading Original Model")
-original_pipeline = pipeline(
-    "image-text-to-text",
-    model=BASE_MODEL_ID,
-    dtype=torch.float16,
-    device_map="auto"
-)
-
-# Load fine-tuned model
-print(f"[INFO] Loading Fine-tuned Model")
-ft_pipe = pipeline(
-    "image-text-to-text",
-    model=FINE_TUNED_MODEL_ID,
-    dtype=torch.bfloat16,
-    device_map="auto"
-)
+# Use bfloat16 for both to avoid CUDA kernel mismatches
+DTYPE = torch.bfloat16 
+
+# Initialize global variables for pipelines
+original_pipeline = None
+ft_pipe = None
+
+def load_models():
+    global original_pipeline, ft_pipe
+    if original_pipeline is None:
+        print(f"[INFO] Loading Original Model")
+        original_pipeline = pipeline(
+            "image-text-to-text",
+            model=BASE_MODEL_ID,
+            torch_dtype=DTYPE,
+            device_map="auto"
+        )
+    if ft_pipe is None:
+        print(f"[INFO] Loading Fine-tuned Model")
+        ft_pipe = pipeline(
+            "image-text-to-text",
+            model=FINE_TUNED_MODEL_ID,
+            torch_dtype=DTYPE,
+            device_map="auto"
+        )
 
 def create_message(input_image):
-    return [{'role': 'user',
- 'content': [{'type': 'image',
-   'image': input_image},
-  {'type': 'text',
-   'text': "Classify the given input image into food or not and if edible food or drink items are present, extract those to a list. If no food/drink items are visible, return empty lists.\n\nOnly return valid JSON in the following form:\n\n```json\n{\n  'is_food': 0, # int - 0 or 1 based on whether food/drinks are present (0 = no foods visible, 1 = foods visible)\n  'image_title': '', # str - short food-related title for what foods/drinks are visible in the image, leave blank if no foods present\n  'food_items': [], # list[str] - list of visible edible food item nouns\n  'drink_items': [] # list[str] - list of visible edible drink item nouns\n}\n```\n"}]}]
+    return [{
+        'role': 'user', 
+        'content': [
+            {'type': 'image', 'image': input_image},
+            {'type': 'text', 'text': "Classify the given input image into food or not and if edible food or drink items are present, extract those to a list. Only return valid JSON."}
+        ]
+    }]
 
 @spaces.GPU
 def extract_foods_from_image(input_image):
+    # Ensure models are loaded on the GPU within the ZeroGPU environment
+    load_models()
+
     input_image = input_image.resize(size=(512, 512))
     input_message = create_message(input_image=input_image)
 
-    # Get outputs from base model (not fine-tuned)
-    original_pipeline_output = original_pipeline(text=[input_message],
-                                                 max_new_tokens=OUTPUT_TOKENS)
-
+    # Note: We pass max_new_tokens and explicitly avoid max_length to prevent the warning/error
+    original_pipeline_output = original_pipeline(
+        text=[input_message], 
+        max_new_tokens=OUTPUT_TOKENS
+    )
     outputs_pretrained = original_pipeline_output[0][0]["generated_text"][-1]["content"]
 
-    # Get outputs from fine-tuned model (fine-tuned on food images)
-    ft_pipe_output = ft_pipe(text=[input_message],
-                             max_new_tokens=OUTPUT_TOKENS)
+    ft_pipe_output = ft_pipe(
+        text=[input_message], 
+        max_new_tokens=OUTPUT_TOKENS
+    )
     outputs_fine_tuned = ft_pipe_output[0][0]["generated_text"][-1]["content"]
 
     return outputs_pretrained, outputs_fine_tuned

requirements.txt

@@ -1,6 +1,7 @@
 num2words
-transformers
+transformers>=4.45.0
 torch
 accelerate
 gradio
 torchvision
+qwen-vl-utils