diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py index 7e3f40883d2..0f6814940eb 100644 --- a/ldm/models/diffusion/ksampler.py +++ b/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) diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index 82839db8751..605cfe9426d 100644 --- a/ldm/simplet2i.py +++ b/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 diff --git a/scripts/dream.py b/scripts/dream.py index 2911e8847ab..0dfc336d7ab 100755 --- a/scripts/dream.py +++ b/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