Update app.py

app.py

@@ -1,7 +1,7 @@
-# app.py for Hugging Face Space: Connecting Meta Llama 3.2 Vision, PaliGemma Segmentation, and Diffusion Model
+# app.py for Hugging Face Space: Connecting Meta Llama 3.2 Vision, Efficient Segmentation, and Diffusion Model
 import gradio as gr
 import spaces  # Import the spaces module to use GPU-specific decorators
-from transformers import PaliGemmaForConditionalGeneration, AutoProcessor, pipeline
+from transformers import VisionEncoderDecoderModel, AutoFeatureExtractor, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 import os
@@ -10,18 +10,18 @@ from PIL import Image
 # Set up Hugging Face token for private model access
 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 = PaliGemmaForConditionalGeneration.from_pretrained(
+# Set up Meta Llama 3.2 Vision model (using Vision Encoder-Decoder model with token)
+llama_vision_model_id = "nlpconnect/vit-gpt2-image-captioning"
+vision_model = VisionEncoderDecoderModel.from_pretrained(
     llama_vision_model_id,
     torch_dtype=torch.bfloat16,
     device_map="auto",
     token=hf_token  # Updated to use 'token' instead of 'use_auth_token'
 )
-processor = AutoProcessor.from_pretrained(llama_vision_model_id, token=hf_token)
+feature_extractor = AutoFeatureExtractor.from_pretrained(llama_vision_model_id, token=hf_token)
 
-# Set up segmentation model using PaliGemma from Hugging Face Hub
-segment_model_id = "google/paligemma-3b-mix-224"
+# Set up segmentation model using an efficient publicly available model
+segment_model_id = "facebook/detr-resnet-50"
 segment_pipe = pipeline(
     "image-segmentation",
     model=segment_model_id,
@@ -40,12 +40,11 @@ diffusion_pipe = diffusion_pipe.to("cuda")  # Force usage of GPU
 @spaces.GPU(duration=120)  # Allocates GPU for a maximum of 120 seconds
 def process_image(image):
     # Step 1: Use Vision model for initial image understanding (captioning)
-    prompt = "<|image|><|begin_of_text|>Describe the image."
-    inputs = processor(image, prompt, return_tensors="pt").to(vision_model.device)
-    output = vision_model.generate(**inputs, max_new_tokens=50)
-    caption = processor.decode(output[0], skip_special_tokens=True)
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(vision_model.device)
+    output_ids = vision_model.generate(pixel_values, max_length=50)
+    caption = vision_model.config.decoder.tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    # Step 2: Segment important parts of the image using PaliGemma
+    # Step 2: Segment important parts of the image using DETR
     segmented_result = segment_pipe(image=image)
     segments = segmented_result