add sanitization for new settings

hydra_node/sanitize.py

@@ -3,7 +3,8 @@ import math
 
 v1pp_defaults = {
     'steps': 50,
-    'plms': True,
+    'sampler': "plms",
+    'image': None,
     'fixed_code': False,
     'ddim_eta': 0.0,
     'height': 512,
@@ -35,6 +36,17 @@ defaults = {
     'dalle-mini': (dalle_mini_defaults, dalle_mini_forced_defaults),
 }
 
+samplers = [
+    "plms",
+    "ddim",
+    "k_euler",
+    "k_euler_ancestral", 
+    "k_heun", 
+    "k_dpm_2", 
+    "k_dpm_2_ancestral", 
+    "k_lms"
+    ]
+
 def closest_multiple(num, mult):
     num_int = int(num)
     floor = math.floor(num_int / mult) * mult
@@ -76,6 +88,9 @@ def sanitize_stable_diffusion(request):
         request.width = closest_multiple(request.width // 2, 64)
         request.height = closest_multiple(request.height // 2, 64)
 
+    if request.sampler not in samplers:
+        return False, "sampler should be one of {}".format(samplers)
+
     return True, request
 
 def sanitize_dalle_mini(request):

main.py

@@ -52,9 +52,10 @@ def root():
 
 class GenerationRequest(BaseModel):
     prompt: str
+    image: str = None
     n_samples: int = 1
     steps: int = 50
-    plms: bool = True
+    sampler: str = "plms"
     fixed_code: bool = False
     ddim_eta: float = 0.0
     height: int = 512