fix: be forgiving of - vs _ in video decoder model name

Addresses #485
brycedrennan · Apr 18, 2024 · efbab3a · efbab3a
1 parent 3a9a397
commit efbab3a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/imaginairy/api/video_sample.py b/imaginairy/api/video_sample.py
@@ -94,6 +94,8 @@ def generate_video(
 
     output_fps = default(output_fps, fps_id)
 
+    model_name = model_name.lower().replace("_", "-")
+
     video_model_config = config.MODEL_WEIGHT_CONFIG_LOOKUP.get(model_name, None)
     if video_model_config is None:
         msg = f"Version {model_name} does not exist."

diff --git a/imaginairy/cli/videogen.py b/imaginairy/cli/videogen.py
@@ -20,7 +20,7 @@
 @click.option(
     "--model",
     default="svd",
-    help="Model to use. One of: svd, svd_xt, svd_image_decoder, svd_xt_image_decoder",
+    help="Model to use. One of: svd, svd-xt, svd-image-decoder, svd-xt-image-decoder",
 )
 @click.option(
     "--fps", default=6, type=int, help="FPS for the AI to target when generating video"