seed for stage 2

hydra_node/models.py

@@ -180,7 +180,7 @@ class StableDiffusionModel(nn.Module):
         if request.seed is not None:
             torch.manual_seed(request.seed)
             np.random.seed(request.seed)
-            
+
         if request.image is not None:
             request.sampler = "ddim_img2img" #enforce ddim for now
             self.ddim.make_schedule(ddim_num_steps=request.steps, ddim_eta=request.ddim_eta, verbose=False)
@@ -316,6 +316,10 @@ class StableDiffusionModel(nn.Module):
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0).squeeze(0)
         x_samples_ddim = pil_upscale(x_samples_ddim, scale=2)
 
+        if request.stage_two_seed is not None:
+            torch.manual_seed(request.stage_two_seed)
+            np.random.seed(request.stage_two_seed)
+
         with torch.autocast("cuda", enabled=self.config.amp):
             with self.model.ema_scope():
                 init_latent = self.model.get_first_stage_encoding(self.model.encode_first_stage(x_samples_ddim))

hydra_node/sanitize.py

@@ -18,6 +18,7 @@ v1pp_defaults = {
     'scale': 7.0,
     'dynamic_threshold': None,
     'seed': None,
+    'stage_two_seed': None,
 }
 
 v1pp_forced_defaults = {

main.py

@@ -71,6 +71,7 @@ class GenerationRequest(BaseModel):
     top_k: int = 256
     grid_size: int = 4
     advanced: bool = False
+    stage_two_seed: int = None
     strength: float = 0.69
     noise: float = 0.667