Update app.py

app.py

@@ -1,6 +1,6 @@
+import os
 import torch
 import gradio as gr
-
 import spaces
 from transformers import pipeline
 
@@ -8,23 +8,66 @@ BASE_MODEL_ID = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
 FINE_TUNED_MODEL_ID = "CreatorJarvis/FoodExtract-Vision-SmolVLM2-500M-fine-tune"
 OUTPUT_TOKENS = 256
 
+DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu"
+if DEVICE_TYPE == "cuda":
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+def _get_dtype(device: str):
+    if device == "cuda":
+        if os.getenv("USE_BF16", "0") == "1":
+            is_bf16_supported = getattr(torch.cuda, "is_bf16_supported", None)
+            if callable(is_bf16_supported) and is_bf16_supported():
+                return torch.bfloat16
+        return torch.float16
+    return torch.float32
+
+DTYPE = _get_dtype(DEVICE_TYPE)
+
+def _make_pipe(model_id: str):
+    device_arg = 0 if DEVICE_TYPE == "cuda" else -1
+    pipe = pipeline(
+        "image-text-to-text",
+        model=model_id,
+        device=device_arg,
+        dtype=DTYPE,
+    )
+    model = getattr(pipe, "model", None)
+    generation_config = getattr(model, "generation_config", None)
+    if generation_config is not None:
+        generation_config.do_sample = False
+        generation_config.max_new_tokens = OUTPUT_TOKENS
+        try:
+            generation_config.max_length = None
+        except Exception:
+            pass
+    return pipe
+
 # 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.bfloat16,
-    device_map="auto"
-)
+original_pipeline = _make_pipe(BASE_MODEL_ID)
 
 # 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"
-)
+ft_pipe = _make_pipe(FINE_TUNED_MODEL_ID)
+
+def _extract_generated_text(pipe_output) -> str:
+    try:
+        item0 = pipe_output[0]
+        if isinstance(item0, dict) and "generated_text" in item0:
+            gt = item0["generated_text"]
+        else:
+            gt = pipe_output[0][0]["generated_text"]
+
+        if isinstance(gt, str):
+            return gt
+        if isinstance(gt, list) and gt:
+            last = gt[-1]
+            if isinstance(last, dict) and "content" in last:
+                return last["content"]
+        return str(gt)
+    except Exception:
+        return str(pipe_output)
 
 def create_message(input_image):
     return [{'role': 'user',
@@ -35,26 +78,28 @@ def create_message(input_image):
 
 @spaces.GPU
 def extract_foods_from_image(input_image):
+    if input_image is None:
+        return "Please upload an image", "Please upload an image"
+
+    input_image = input_image.convert("RGB")
     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)
+    original_pipeline_output = original_pipeline(text=[input_message])
 
-    outputs_pretrained = original_pipeline_output[0][0]["generated_text"][-1]["content"]
+    outputs_pretrained = _extract_generated_text(original_pipeline_output)
 
     # Get outputs from fine-tuned model (fine-tuned on food images)
-    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"]
+    ft_pipe_output = ft_pipe(text=[input_message])
+    outputs_fine_tuned = _extract_generated_text(ft_pipe_output)
 
     return outputs_pretrained, outputs_fine_tuned
 
 demo_title = "🥑➡️📝 FoodExtract-Vision with a fine-tuned SmolVLM2-500M"
 demo_description = """* **Base model:** https://huggingface.co/HuggingFaceTB/SmolVLM-500M-Instruct
 * **Fine-tuning dataset:** https://huggingface.co/datasets/mrdbourke/FoodExtract-1k-Vision (1k food images and 500 not food images)
-* **Fine-tuned model:** https://huggingface.co/mrdbourke/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1
+* **Fine-tuned model:** https://huggingface.co/CreatorJarvis/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1
 
 ## Overview
 
@@ -96,9 +141,7 @@ demo = gr.Interface(
     description=demo_description,
     outputs=[gr.Textbox(lines=4, label="Original Model (not fine-tuned)"),
              gr.Textbox(lines=4, label="Fine-tuned Model")],
-    examples=[["examples/camera.jpeg"],
-              ["examples/Tandoori-Chicken.jpg"],
-              ["examples/fries.jpeg"]],
+    
 )
 
 if __name__ == "__main__":