Update app.py

app.py

@@ -20,10 +20,14 @@ vision_model = MllamaForConditionalGeneration.from_pretrained(
 )
 processor = AutoProcessor.from_pretrained(llama_vision_model_id, token=hf_token)
 
-# Set up segmentation model using Segment Anything 2 (sam2_hiera_small.pt)
+# Set up segmentation model using Segment Anything 2 from Hugging Face Hub
 segment_model_id = "camenduru/segment-anything-2"
-segment_model_path = "sam2_hiera_small.pt"
-segment_pipe = torch.load(segment_model_path, map_location="cuda")  # Load segmentation model on GPU
+segment_pipe = pipeline(
+    "image-segmentation",
+    model=segment_model_id,
+    device=0,  # Force usage of GPU
+    token=hf_token  # Updated to use 'token'
+)
 
 # Set up Stable Diffusion Lite model
 stable_diffusion_model_id = "runwayml/stable-diffusion-v1-5"
@@ -42,8 +46,7 @@ def process_image(image):
     caption = processor.decode(output[0], skip_special_tokens=True)
 
     # Step 2: Segment important parts of the image using Segment Anything 2
-    # Use the loaded segment model to perform segmentation
-    segmented_result = segment_pipe(image=image)  # Assuming a callable model or appropriate method
+    segmented_result = segment_pipe(image=image)
     segments = segmented_result
 
     # Step 3: Modify segmented image using Diffusion model