Update app_recon.py

app_recon.py

@@ -55,6 +55,25 @@ from torchvision.transforms import InterpolationMode
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
 
+# Initialization.. Downloading the ckpts
+
+image_encoder_checkpoint1 = os.path.join(os.path.dirname(__file__), "image_encoder", "model.safetensors")
+image_encoder_checkpoint2 = os.path.join(os.path.dirname(__file__), "image_encoder", "pytorch_model.bin")
+vae_checkpoint1 = os.path.join(os.path.dirname(__file__), "vae", "diffusion_pytorch_model.bin")
+vae_checkpoint2 = os.path.join(os.path.dirname(__file__), "vae", "diffusion_pytorch_model.safetensors")
+ckpt_lists = [image_encoder_checkpoint1, image_encoder_checkpoint2, vae_checkpoint1, vae_checkpoint2]
+
+image_encoder_url1 = 'https://huggingface.co/lemonaddie/Geowizard/resolve/main/image_encoder/model.safetensors'
+image_encoder_url2 = 'https://huggingface.co/lemonaddie/Geowizard/resolve/main/image_encoder/pytorch_model.bin'
+vae_url1 = 'https://huggingface.co/lemonaddie/Geowizard/resolve/main/vae/diffusion_pytorch_model.bin'
+vae_url2 = 'https://huggingface.co/lemonaddie/Geowizard/resolve/main/vae/diffusion_pytorch_model.safetensors'
+url_lists = [image_encoder_url1, image_encoder_url2, vae_url1, vae_url2]
+
+for ckpt_path, ckpt_url in zip(ckpt_lists, url_lists):
+    if not os.path.exists(ckpt_path):
+        print("Downloading to " + ckpt_path + "...")
+        os.system('wget -P' + os.path.dirname(ckpt_path) + ' -nv ' + ckpt_url)
+
 vae = AutoencoderKL.from_pretrained("./", subfolder='vae')
 scheduler = DDIMScheduler.from_pretrained("./", subfolder='scheduler')
 image_encoder = CLIPVisionModelWithProjection.from_pretrained("./", subfolder="image_encoder")
@@ -96,7 +115,7 @@ def sam_init():
     sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_vit_h_4b8939.pth")
 
     if not os.path.exists(sam_checkpoint):
-        print("Downloading SAM-H checkpoint to " + os.path.exists(sam_checkpoint))
+        print("Downloading SAM-H checkpoint to " + sam_checkpoint)
         os.system('wget -P' + os.path.dirname(__file__) + ' -nv https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth')
     
     model_type = "vit_h"