Masker Dict

main.py

@@ -6,7 +6,7 @@ from sentry_sdk import capture_exception
 from sentry_sdk import capture_message
 from sentry_sdk import start_transaction
 from hydra_node.config import init_config_model
-from typing import Optional, List
+from typing import Optional, List, TypedDict
 import socket
 from hydra_node.sanitize import sanitize_input
 import uvicorn
@@ -54,6 +54,10 @@ def shutdown_event():
 def root():
     return "OK"
 
+class Masker(TypedDict):
+    seed: int
+    mask: str
+
 class GenerationRequest(BaseModel):
     prompt: str
     image: str = None
@@ -78,7 +82,7 @@ class GenerationRequest(BaseModel):
     noise: float = 0.667
     mitigate: bool = False
     module: str = None
-    masks: List[str] = None
+    masks: List[Masker] = None
 
 class GenerationOutput(BaseModel):
     output: List[str]