bit hacky for now

hydra_node/sanitize.py

@@ -34,19 +34,7 @@ defaults = {
     'dalle-mini': (dalle_mini_defaults, dalle_mini_forced_defaults),
 }
 
-def sanitize_input(config, request):
-    """
-    Sanitize the input data and set defaults
-    """
-    request = DotMap(request)
-    default, forced_default = defaults[config.model]
-    for k, v in default.items():
-        if k not in request:
-            request[k] = v
-
-    for k, v in forced_default.items():
-        request[k] = v
-
+def sanitize_stable_diffusion(request):
     if request.width * request.height == 0:
         return False, "width and height must be non-zero"
 
@@ -73,6 +61,26 @@ def sanitize_input(config, request):
 
     return True, request
 
+def sanitize_dalle_mini(request):
+    return True, request
+
+def sanitize_input(config, request):
+    """
+    Sanitize the input data and set defaults
+    """
+    request = DotMap(request)
+    default, forced_default = defaults[config.model]
+    for k, v in default.items():
+        if k not in request:
+            request[k] = v
+
+    for k, v in forced_default.items():
+        request[k] = v
+
+    if config.model == 'stable-diffusion':
+        return sanitize_stable_diffusion(request)
+    elif config.model == 'dalle-mini':
+        return sanitize_dalle_mini(request)