Update app.py

app.py

@@ -20,14 +20,10 @@ vision_model = MllamaForConditionalGeneration.from_pretrained(
 )
 processor = AutoProcessor.from_pretrained(llama_vision_model_id, token=hf_token)
 
-# Set up Meta Segment Anything 2 model (using private model with token)
-segment_model_id = "meta/segment-anything-2"
-segment_pipe = pipeline(
-    "image-segmentation",
-    model=segment_model_id,
-    device=0,  # Force usage of GPU
-    token=hf_token  # Updated to use 'token' instead of 'use_auth_token'
-)
+# Set up segmentation model using Segment Anything 2 (sam2_hiera_small.pt)
+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
 
 # Set up Stable Diffusion Lite model
 stable_diffusion_model_id = "runwayml/stable-diffusion-v1-5"
@@ -45,9 +41,10 @@ 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
-    segmented_result = segment_pipe(image=image)
-    segments = segmented_result["segments"]
+    # 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
+    segments = segmented_result
 
     # Step 3: Modify segmented image using Diffusion model
     # Here, we modify based on the caption result and segmented area