Create dalle_models.py

dalle_models.py

@@ -0,0 +1,24 @@
+from enum import Enum 
+
+IMAGES_OUTPUT_DIR = 'generations'
+
+DALLE_MODEL_MINI = "dalle-mini/dalle-mini/mini-1:v0"  # the original DALL-E Mini. Fastest yet suboptimal results
+DALLE_MODEL_MEGA = "dalle-mini/dalle-mini/mega-1-fp16:latest"  # the advanced version of DALL-E Mini. Requires more compute and VRAM
+DALLE_MODEL_MEGA_FULL = "dalle-mini/dalle-mini/mega-1:latest"  # DALL-E Mega. Warning: requires significantly more storage and GPU RAM
+DALLE_COMMIT_ID = None
+
+# VQGAN model
+VQGAN_REPO = "dalle-mini/vqgan_imagenet_f16_16384"
+VQGAN_COMMIT_ID = "e93a26e7707683d349bf5d5c41c5b0ef69b677a9"
+
+
+# We can customize generation parameters (see https://huggingface.co/blog/how-to-generate)
+GEN_TOP_K = None
+GEN_TOP_P = None
+TEMPERATURE  = None
+COND_SCALE = 10.0
+
+class ModelSize(Enum):
+    MINI = "Mini"
+    MEGA = "Mega"
+    MEGA_FULL = "Mega_full"