Update app.py

app.py

@@ -163,54 +163,10 @@ 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 = build_sam(checkpoint='sam_vit_h_4b8939.pth')
+        self.sam = sam_model_registry["vit_h"]()
+        self.sam.load_state_dict(torch.utils.model_zoo.load_url(MODEL_DICT["vit_h"]))
         self.sam.to(device=self.device)
         self.sam_predictor = SamPredictor(self.sam)