diff --git a/cog.yaml b/cog.yaml
index 5317123..f8b1422 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -67,10 +67,10 @@ build:
     - "wandb==0.15.11"
     - "zipp==3.17.0"
     - "xformers"
+    - "git+https://github.com/hotshotco/Hotshot-XL"
 
   run:
     - apt-get update && apt-get install -y git-lfs ffmpeg
-    - git clone https://github.com/hotshotco/Hotshot-XL /Hotshot-XL
 
 
 # predict.py defines how predictions are run on your model
diff --git a/scripts/download_weights.py b/scripts/download_weights.py
index c1ef274..1e20bf2 100755
--- a/scripts/download_weights.py
+++ b/scripts/download_weights.py
@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
-from diffusers import DiffusionPipeline
+from hotshot_xl.pipelines.hotshot_xl_pipeline import HotshotXLPipeline
 import torch
 
-pipe = DiffusionPipeline.from_pretrained(
-    "hotshotco/SDXL-512",
+pipe = HotshotXLPipeline.from_pretrained(
+    "hotshotco/Hotshot-XL",
     torch_dtype=torch.float16,
     use_safetensors=True
 )