vanilla check

hydra_node/models.py

@@ -213,8 +213,12 @@ class StableDiffusionModel(nn.Module):
     @torch.autocast("cuda", enabled=True, dtype=torch.float16)
     def sample(self, request):
         if request.module is not None:
-            module = self.premodules[request.module]
-            CrossAttention.set_hypernetwork(module)
+            if request.module == "vanilla":
+                pass
+            
+            else:
+                module = self.premodules[request.module]
+                CrossAttention.set_hypernetwork(module)
 
         if request.seed is not None:
             torch.manual_seed(request.seed)

main.py

@@ -78,6 +78,7 @@ class GenerationRequest(BaseModel):
     noise: float = 0.667
     mitigate: bool = False
     module: str = None
+    masks: List[str] = None
 
 class GenerationOutput(BaseModel):
     output: List[str]