add uc

hydra_node/models.py

@@ -354,7 +354,11 @@ class StableDiffusionModel(nn.Module):
 
         uc = None
         if request.scale != 1.0:
-            uc = self.model.get_learned_conditioning(request.n_samples * [""])
+            if request.uc is not None:
+                uc = [request.uc] * request.n_samples
+                uc = prompt_mixing(self.model, uc[0], request.n_samples)
+            else:
+                uc = self.model.get_learned_conditioning(request.n_samples * [""])
 
         shape = [
             request.latent_channels,

main.py

@@ -98,6 +98,7 @@ class GenerationRequest(BaseModel):
     mitigate: bool = False
     module: str = None
     masks: List[Masker] = None
+    uc: str = None
 
 class TextRequest(BaseModel):
     prompt: str