update app.py

app.py

@@ -1,9 +1,12 @@
+import os
 import cv2
-import numpy as np
-import torch
 import time
-from PIL import Image, ImageDraw
+import torch
+import subprocess
+import numpy as np
 import gradio as gr
+import urllib.request
+from PIL import Image, ImageDraw
 import matplotlib.pyplot as plt
 
 from Garage.models.GroundedSegmentAnything.segment_anything.segment_anything import SamPredictor, sam_model_registry
@@ -43,6 +46,7 @@ class GradioWindow():
         self.SAM_CHECKPOINT_PATH = SAM_CHECKPOINT_PATH
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.download_weights()
         # for debug
         # self.augmenter = None
         self.augmenter = Augmenter(device=self.device)
@@ -159,6 +163,52 @@ class GradioWindow():
                 outputs=[augmented_img, generated_prompt],
             )
 
+    def download_weights(self):
+        models = [
+            "https://huggingface.co/JunhaoZhuang/PowerPaint-v2-1/",
+            "https://huggingface.co/llava-hf/llava-1.5-7b-hf",
+            "https://huggingface.co/danulkin/llama",
+        ]
+
+        destinations = [
+            "Garage/models/checkpoints/ppt-v2-1",
+            "Garage/models/checkpoints/llava-1.5-7b-hf",
+            "Garage/models/checkpoints/llama-3-8b-Instruct",
+        ]
+
+        if not os.path.exists("Garage/models/checkpoints"):
+            os.makedirs("Garage/models/checkpoints")
+
+        for model, destination in zip(models, destinations):
+            # Git LFS clone command
+            command = ["git", "lfs", "clone", model, destination]
+            try:
+                result = subprocess.run(command, check=True, text=True, capture_output=True)
+                print("Command Output:", result.stdout)
+
+            except subprocess.CalledProcessError as e:
+                print(f"Error: {e}")
+                print("Command Output:", e.output)
+
+        models = [
+            "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+            "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth",
+        ]
+
+        destinations = [
+            "Garage/models/checkpoints/GroundedSegmentAnything/sam_vit_h_4b8939.pth",
+            "Garage/models/checkpoints/GroundedSegmentAnything/groundingdino_swint_ogc.pth",
+        ]
+        if not os.path.exists("Garage/models/checkpoints/GroundedSegmentAnything"):
+            os.makedirs("Garage/models/checkpoints/GroundedSegmentAnything")
+
+        for model, destination in zip(models, destinations):
+            if not os.path.exists(destination):
+                urllib.request.urlretrieve(model, destination)
+                print(f"Downloaded {model} to {destination}")
+            else:    
+                print(f"Model {model} already exists")
+
     def setup_model(self) -> SamPredictor:
         self.sam = sam_model_registry[self.model_type](checkpoint=self.SAM_CHECKPOINT_PATH)
         self.sam.to(device=self.device)