seed fuzzing

ldm/models/diffusion/ksampler.py

@@ -66,8 +66,8 @@ def route_callback(k_callback_values):
                 img_callback(k_callback_values['x'], k_callback_values['i'])
 
         sigmas = self.model.get_sigmas(S)
-        if x_T:
-            x = x_T
+        if x_T is not None:
+            x = x_T * sigmas[0]
         else:
             x = (
                 torch.randn([batch_size, *shape], device=self.device)

ldm/simplet2i.py

@@ -214,6 +214,8 @@ def prompt2image(
         variants=None,
         sampler_name=None,
         log_tokenization=False,
+        variant_amount=0.0,
+        variant_seed=None,
         **args,
     ):   # eat up additional cruft
         """
@@ -233,6 +235,8 @@ def prompt2image(
            variants                        // if >0, the 1st generated image will be passed back to img2img to generate the requested number of variants
            step_callback                   // a function or method that will be called each step
            image_callback                  // a function or method that will be called each time an image is generated
+           variant_amount                  // optional 0-1 value to slerp from -S noise to random noise (allows variations on an image)
+           variant_seed                    // optional target seed that -S noise is slerped to (interpolate one image to another)
 
         To use the step callback, define a function that receives two arguments:
         - Image GPU data
@@ -266,6 +270,9 @@ def process_image(image,seed):
         assert (
             0.0 <= strength <= 1.0
         ), 'can only work with strength in [0.0, 1.0]'
+        assert (
+            0.0 <= variant_amount <= 1.0
+        ), '-v --variant_amount must be in 0.0 to 1.0 range'
         w, h = map(
             lambda x: x - x % 64, (width, height)
         )  # resize to integer multiple of 64
@@ -314,12 +321,19 @@ def process_image(image,seed):
                     width=width,
                     height=height,
                     callback=step_callback,
+                    variant_amount=variant_amount,
+                    variant_seed=variant_seed,
                 )
 
             with scope(self.device.type), self.model.ema_scope():
                 for n in trange(iterations, desc='Generating'):
                     seed_everything(seed)
-                    image = next(images_iterator)
+                    image,modified_seed = next(images_iterator)
+                    # images_iterator can modify seed when using -v  
+                    # it returns it to be used in filename/log
+                    if modified_seed is not None:
+                        seed = modified_seed
+
                     results.append([image, seed])
                     if image_callback is not None:
                         image_callback(image, seed)
@@ -396,12 +410,16 @@ def _txt2img(
         width,
         height,
         callback,
+        variant_amount,
+        variant_seed,
     ):
         """
         An infinite iterator of images from the prompt.
         """
 
         sampler = self.sampler
+
+        base_x_T = self._get_base_noise(width, height, variant_amount)
 
         while True:
             uc, c = self._get_uc_and_c(prompt, skip_normalize)
@@ -410,6 +428,9 @@ def _txt2img(
                 height // self.downsampling_factor,
                 width // self.downsampling_factor,
             ]
+
+            x_T,seed = self._apply_variation_slerp(width, height, variant_amount, variant_seed, base_x_T)
+
             samples, _ = sampler.sample(
                 batch_size=1,
                 S=steps,
@@ -419,9 +440,10 @@ def _txt2img(
                 unconditional_guidance_scale=cfg_scale,
                 unconditional_conditioning=uc,
                 eta=ddim_eta,
-                img_callback=callback
+                img_callback=callback,
+                x_T = x_T
             )
-            yield self._sample_to_image(samples)
+            yield self._sample_to_image(samples),seed
 
     @torch.no_grad()
     def _img2img(
@@ -480,7 +502,7 @@ def _img2img(
                 unconditional_guidance_scale=cfg_scale,
                 unconditional_conditioning=uc,
             )
-            yield self._sample_to_image(samples)
+            yield self._sample_to_image(samples),None #none is to stay consistent with _txt2img
 
     # TODO: does this actually need to run every loop? does anything in it vary by random seed?
     def _get_uc_and_c(self, prompt, skip_normalize):
@@ -667,3 +689,79 @@ def _log_tokenization(self, text):
         print(f"\nTokens ({usedTokens}):\n{tokenized}\x1b[0m")
         if discarded != "":
             print(f"Tokens Discarded ({totalTokens-usedTokens}):\n{discarded}\x1b[0m")
+
+    def _get_base_noise(self, width:int, height:int, variant_amount:float) -> torch.Tensor:
+        base_x_T = None
+        if variant_amount != 0.0:
+            variant_amount = max(0.0, min(1.0, variant_amount))
+            # base noise is made from whatever our seed currently is
+            base_x_T = torch.randn([1,
+                                    self.latent_channels,
+                                    height // self.downsampling_factor,
+                                    width  // self.downsampling_factor],
+                                    device=self.device)
+        return base_x_T
+
+    def _apply_variation_slerp(self, 
+        width:int, height:int,
+        variant_amount:float, variant_seed:int, 
+        base_x_T:torch.Tensor) -> "[torch.Tensor,int]":
+        x_T = None
+        seed = None
+        if variant_amount != 0.0:
+            variant_amount = max(0.0, min(1.0, variant_amount))
+            # no variant seed specified, generate random noise
+            if variant_seed is None:
+                random.seed() # reset RNG to an actually random state, so we can get a random seed
+                seed = random.randrange(0,np.iinfo(np.uint32).max)
+            else: # use variant seed for noise
+                seed = variant_seed
+
+            seed_everything(seed)
+
+            target_x_T = torch.randn([1,
+                self.latent_channels,
+                height // self.downsampling_factor,
+                width  // self.downsampling_factor],
+                device=self.device)
+
+            # slerp base -> target using variant amount
+            x_T = self.slerp(variant_amount, base_x_T, target_x_T)
+        return x_T, seed
+
+    def slerp(self, t, v0, v1, DOT_THRESHOLD=0.9995):
+        '''
+        Spherical linear interpolation
+        Args:
+            t (float/np.ndarray): Float value between 0.0 and 1.0
+            v0 (np.ndarray): Starting vector
+            v1 (np.ndarray): Final vector
+            DOT_THRESHOLD (float): Threshold for considering the two vectors as
+                                colineal. Not recommended to alter this.
+        Returns:
+            v2 (np.ndarray): Interpolation vector between v0 and v1
+        '''
+        inputs_are_torch = False
+        if not isinstance(v0,np.ndarray):
+            inputs_are_torch = True
+            v0 = v0.detach().cpu().numpy()
+        if not isinstance(v1,np.ndarray):
+            inputs_are_torch = True
+            v1 = v1.detach().cpu().numpy()
+
+        dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
+        if np.abs(dot) > DOT_THRESHOLD:
+            v2 = (1 - t) * v0 + t * v1
+        else:
+            theta_0 = np.arccos(dot)
+            sin_theta_0 = np.sin(theta_0)
+            theta_t = theta_0 * t
+            sin_theta_t = np.sin(theta_t)
+            s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+            s1 = sin_theta_t / sin_theta_0
+            v2 = s0 * v0 + s1 * v1
+
+        if inputs_are_torch:
+            v2 = torch.from_numpy(v2).to(self.device)
+
+        return v2

scripts/dream.py

@@ -536,6 +536,20 @@ def create_cmd_parser():
         action='store_true',
         help='shows how the prompt is split into tokens'
     )
+    parser.add_argument(
+        '-v',
+        '--variant_amount',
+        default=0.0,
+        type=float,
+        help='0.0 to 1.0 value controlling percentage of variation noise applied'
+    )
+    parser.add_argument(
+        '-V',
+        '--variant_seed',
+        default=None,
+        type=int,
+        help='manually seed the variation noise, instead of using randomly generated noise'
+    )
     return parser