Update app.py

app.py

@@ -1,7 +1,7 @@
-# app.py for Hugging Face Space: Connecting Meta Llama 3.2 Vision, Segment Anything 2, and Diffusion Model
+# app.py for Hugging Face Space: Connecting Meta Llama 3.2 Vision, PaliGemma Segmentation, and Diffusion Model
 import gradio as gr
 import spaces  # Import the spaces module to use GPU-specific decorators
-from transformers import MllamaForConditionalGeneration, AutoProcessor, pipeline
+from transformers import PaliGemmaForConditionalGeneration, AutoProcessor, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 import os
@@ -12,7 +12,7 @@ hf_token = os.getenv("HF_TOKEN")  # Fetch token from repository secrets
 
 # Set up Meta Llama 3.2 Vision model (using private model with token)
 llama_vision_model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-vision_model = MllamaForConditionalGeneration.from_pretrained(
+vision_model = PaliGemmaForConditionalGeneration.from_pretrained(
     llama_vision_model_id,
     torch_dtype=torch.bfloat16,
     device_map="auto",
@@ -20,8 +20,8 @@ vision_model = MllamaForConditionalGeneration.from_pretrained(
 )
 processor = AutoProcessor.from_pretrained(llama_vision_model_id, token=hf_token)
 
-# Set up segmentation model using MaskFormer Swin Large from Hugging Face Hub
-segment_model_id = "facebook/maskformer-swin-large"
+# Set up segmentation model using PaliGemma from Hugging Face Hub
+segment_model_id = "google/paligemma-3b-mix-224"
 segment_pipe = pipeline(
     "image-segmentation",
     model=segment_model_id,
@@ -45,7 +45,7 @@ def process_image(image):
     output = vision_model.generate(**inputs, max_new_tokens=50)
     caption = processor.decode(output[0], skip_special_tokens=True)
 
-    # Step 2: Segment important parts of the image using MaskFormer Swin Large
+    # Step 2: Segment important parts of the image using PaliGemma
     segmented_result = segment_pipe(image=image)
     segments = segmented_result