Fix fetching the wrong model, moved hotshot to py deps

lucataco · Oct 11, 2023 · 64431b9 · 64431b9
1 parent 3222223
commit 64431b9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git 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
@@ -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
 )