update

black_box_image_edit/cosxl_edit.py

@@ -21,11 +21,17 @@ class CosXLEdit():
         from .cosxl.utils import set_timesteps_patched
 
         EDMEulerScheduler.set_timesteps = set_timesteps_patched
-        #edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
-        edit_file_path = "./black_box_image_edit/cosxl/cosxl_edit.safetensors"
-        self.pipe = CosStableDiffusionXLInstructPix2PixPipeline.from_single_file(
-            edit_file_path, num_in_channels=8
-        )
+
+        try:
+            edit_file = hf_hub_download(repo_id="TIGER-Lab/cosxl", filename="cosxl_edit.safetensors")
+            self.pipe = CosStableDiffusionXLInstructPix2PixPipeline.from_single_file(
+                edit_file, num_in_channels=8
+            )
+        except:
+            edit_file_path = "./black_box_image_edit/cosxl/cosxl_edit.safetensors"
+            self.pipe = CosStableDiffusionXLInstructPix2PixPipeline.from_single_file(
+                edit_file_path, num_in_channels=8
+            )
         self.pipe.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
         self.pipe.to(device)
         self.pipe.enable_vae_tiling()