Add support for v2.1

castorini · Dec 16, 2022 · 1f99576 · 1f99576
1 parent 280231c
commit 1f99576
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/daam/run/generate.py b/daam/run/generate.py
@@ -23,7 +23,9 @@ def main():
     model_id_map = {
         'v1': 'runwayml/stable-diffusion-v1-5',
         'v2-base': 'stabilityai/stable-diffusion-2-base',
-        'v2-large': 'stabilityai/stable-diffusion-2'
+        'v2-large': 'stabilityai/stable-diffusion-2',
+        'v2-1-base': 'stabilityai/stable-diffusion-2-1-base',
+        'v2-1-large': 'stabilityai/stable-diffusion-2-1'
     }
 
     parser = argparse.ArgumentParser()