From 0febcc2aa84131d2ec737056eabb0ad1ad35fe29 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 5 Dec 2023 14:07:52 -0500 Subject: [PATCH] animatediff full latent mode --- CHANGELOG.md | 2 + README.md | 3 +- modules/processing_correction.py | 76 ++++++++++++++++++++------------ modules/processing_diffusers.py | 24 +++++++--- scripts/animatediff.py | 2 +- 5 files changed, 70 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 481507267..b1c40ef95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ - **Diffusers** - **IP Adapter** add support for `ip-adapter-plus_sd15` and `ip-adapter-plus-face_sd15` + - **AnimateDiff** can now be used with *second pass* and *hdr controls* - enhance, upscale and hires your videos! + - **HDE controls** are now batch-aware for enhancement of multiple images - added support for basic [ModelScope T2V](https://huggingface.co/damo-vilab/text-to-video-ms-1.7b) model - simply select from *networks -> reference* and use from *txt2img* tab - **General** diff --git a/README.md b/README.md index 2f22500ca..f2841e787 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,7 @@ Additional models will be added as they become available and there is public int - [RunwayML Stable Diffusion](https://github.com/Stability-AI/stablediffusion/) 1.x and 2.x *(all variants)* - [StabilityAI Stable Diffusion XL](https://github.com/Stability-AI/generative-models) -- [StabilityAI Stable Video Diffusion Base and XT](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid) +- [StabilityAI Stable Video Diffusion](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid) Base and XT - [Segmind SSD-1B](https://huggingface.co/segmind/SSD-1B) - [LCM: Latent Consistency Models](https://github.com/openai/consistency_models) - [Kandinsky](https://github.com/ai-forever/Kandinsky-2) *2.1 and 2.2 and latest 3.0* @@ -68,6 +68,7 @@ Additional models will be added as they become available and there is public int - [Warp Wuerstchen](https://huggingface.co/blog/wuertschen) - [Tsinghua UniDiffusion](https://github.com/thu-ml/unidiffuser) - [DeepFloyd IF](https://github.com/deep-floyd/IF) *Medium and Large* +- [ModelScope T2V](https://huggingface.co/damo-vilab/text-to-video-ms-1.7b) - [Segmind SD Distilled](https://huggingface.co/blog/sd_distillation) *(all variants)* Also supported are modifiers such as: diff --git a/modules/processing_correction.py b/modules/processing_correction.py index 3e8d6994e..d0f189dcc 100644 --- a/modules/processing_correction.py +++ b/modules/processing_correction.py @@ -11,56 +11,74 @@ from modules import shared debug = shared.log.info if os.environ.get('SD_HDR_DEBUG', None) is not None else lambda *args, **kwargs: None -def soft_clamp_tensor(input_tensor, threshold=0.8, boundary=4): +def soft_clamp_tensor(tensor, threshold=0.8, boundary=4): # shrinking towards the mean; will also remove outliers - if max(abs(input_tensor.max()), abs(input_tensor.min())) < boundary or threshold == 0: - return input_tensor + if max(abs(tensor.max()), abs(tensor.min())) < boundary or threshold == 0: + return tensor channel_dim = 1 threshold *= boundary - max_vals = input_tensor.max(channel_dim, keepdim=True)[0] - max_replace = ((input_tensor - threshold) / (max_vals - threshold)) * (boundary - threshold) + threshold - over_mask = input_tensor > threshold - min_vals = input_tensor.min(channel_dim, keepdim=True)[0] - min_replace = ((input_tensor + threshold) / (min_vals + threshold)) * (-boundary + threshold) - threshold - under_mask = input_tensor < -threshold + max_vals = tensor.max(channel_dim, keepdim=True)[0] + max_replace = ((tensor - threshold) / (max_vals - threshold)) * (boundary - threshold) + threshold + over_mask = tensor > threshold + min_vals = tensor.min(channel_dim, keepdim=True)[0] + min_replace = ((tensor + threshold) / (min_vals + threshold)) * (-boundary + threshold) - threshold + under_mask = tensor < -threshold debug(f'HDE soft clamp: threshold={threshold} boundary={boundary}') - input_tensor = torch.where(over_mask, max_replace, torch.where(under_mask, min_replace, input_tensor)) - return input_tensor + tensor = torch.where(over_mask, max_replace, torch.where(under_mask, min_replace, tensor)) + return tensor -def center_tensor(input_tensor, channel_shift=1.0, full_shift=1.0, channels=[0, 1, 2, 3]): # pylint: disable=dangerous-default-value # noqa: B006 +def center_tensor(tensor, channel_shift=1.0, full_shift=1.0, channels=[0, 1, 2, 3]): # pylint: disable=dangerous-default-value # noqa: B006 if channel_shift == 0 and full_shift == 0: - return input_tensor + return tensor means = [] for channel in channels: - means.append(input_tensor[0, channel].mean()) - input_tensor[0, channel] -= means[-1] * channel_shift + means.append(tensor[0, channel].mean()) + tensor[0, channel] -= means[-1] * channel_shift debug(f'HDR center: channel-shift{channel_shift} full-shift={full_shift} means={torch.stack(means)}') - input_tensor = input_tensor - input_tensor.mean() * full_shift - return input_tensor + tensor = tensor - tensor.mean() * full_shift + return tensor -def maximize_tensor(input_tensor, boundary=1.0, channels=[0, 1, 2]): # pylint: disable=dangerous-default-value # noqa: B006 +def maximize_tensor(tensor, boundary=1.0, channels=[0, 1, 2]): # pylint: disable=dangerous-default-value # noqa: B006 if boundary == 1.0: - return input_tensor + return tensor boundary *= 4 - min_val = input_tensor.min() - max_val = input_tensor.max() + min_val = tensor.min() + max_val = tensor.max() normalization_factor = boundary / max(abs(min_val), abs(max_val)) - input_tensor[0, channels] *= normalization_factor + tensor[0, channels] *= normalization_factor debug(f'HDR maximize: boundary={boundary} min={min_val} max={max_val} factor={normalization_factor}') - return input_tensor + return tensor -def correction_callback(p, timestep, kwags): +def correction(p, timestep, latent): if timestep > 950 and p.hdr_clamp: p.extra_generation_params["HDR clamp"] = f'{p.hdr_threshold}/{p.hdr_boundary}' - kwags["latents"] = soft_clamp_tensor(kwags["latents"], threshold=p.hdr_threshold, boundary=p.hdr_boundary) + latent = soft_clamp_tensor(latent, threshold=p.hdr_threshold, boundary=p.hdr_boundary) if timestep > 700 and p.hdr_center: p.extra_generation_params["HDR center"] = f'{p.hdr_channel_shift}/{p.hdr_full_shift}' - kwags["latents"] = center_tensor(kwags["latents"], channel_shift=p.hdr_channel_shift, full_shift=p.hdr_full_shift) + latent = center_tensor(latent, channel_shift=p.hdr_channel_shift, full_shift=p.hdr_full_shift) if timestep > 1 and timestep < 100 and p.hdr_maximize: p.extra_generation_params["HDR max"] = f'{p.hdr_max_center}/p.hdr_max_boundry' - kwags["latents"] = center_tensor(kwags["latents"], channel_shift=p.hdr_max_center, full_shift=1.0) - kwags["latents"] = maximize_tensor(kwags["latents"], boundary=p.hdr_max_boundry) - return kwags + latent = center_tensor(latent, channel_shift=p.hdr_max_center, full_shift=1.0) + latent = maximize_tensor(latent, boundary=p.hdr_max_boundry) + return latent + + +def correction_callback(p, timestep, kwargs): + if not p.hdr_clamp and not p.hdr_center and not p.hdr_maximize: + return kwargs + latents = kwargs["latents"] + if len(latents.shape) == 4: # standard batched latent + for i in range(latents.shape[0]): + latents[i] = correction(p, timestep, latents[i]) + elif len(latents.shape) == 5 and latents.shape[0] == 1: # probably animatediff + latents = latents.squeeze(0).permute(1, 0, 2, 3) + for i in range(latents.shape[0]): + latents[i] = correction(p, timestep, latents[i]) + latents = latents.permute(1, 0, 2, 3).unsqueeze(0) + else: + shared.log.debug(f'HDR correction: unknown latent shape {latents.shape}') + kwargs["latents"] = latents + return kwargs diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 95b2c1c41..512b67314 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -37,6 +37,10 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro p.mask_for_overlay = images.resize_image(1, p.mask_for_overlay, tgt_width, tgt_height, upscaler_name=None) def hires_resize(latents): # input=latents output=pil + if not torch.is_tensor(latents): + shared.log.warning('Hires: input is not tensor') + first_pass_images = vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil') + return first_pass_images latent_upscaler = shared.latent_upscale_modes.get(p.hr_upscaler, None) shared.log.info(f'Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}') if latent_upscaler is not None: @@ -59,9 +63,10 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro for j in range(len(decoded)): images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix) - def diffusers_callback_legacy(step: int, _timestep: int, latents: torch.FloatTensor): + def diffusers_callback_legacy(step: int, timestep: int, latents: torch.FloatTensor): shared.state.sampling_step = step shared.state.current_latent = latents + latents = correction_callback(p, timestep, {'latents': latents}) if shared.state.interrupted or shared.state.skipped: raise AssertionError('Interrupted...') if shared.state.paused: @@ -386,9 +391,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro shared.log.debug(f'Diffuser pipeline: {model.__class__.__name__} task={sd_models.get_diffusers_task(model)} set={clean}') if p.hdr_clamp or p.hdr_center or p.hdr_maximize: txt = 'HDR:' - txt += f' Clamp threshold={p.hdr_threshold} boundary={p.hdr_boundary}' if p.hdr_clamp else 'Clamp off' - txt += f' Center channel-shift={p.hdr_channel_shift} full-shift={p.hdr_full_shift}' if p.hdr_center else 'Center off' - txt += f' Maximize boundary={p.hdr_max_boundry} center={p.hdr_max_center}' if p.hdr_maximize else 'Maximize off' + txt += f' Clamp threshold={p.hdr_threshold} boundary={p.hdr_boundary}' if p.hdr_clamp else ' Clamp off' + txt += f' Center channel-shift={p.hdr_channel_shift} full-shift={p.hdr_full_shift}' if p.hdr_center else ' Center off' + txt += f' Maximize boundary={p.hdr_max_boundry} center={p.hdr_max_center}' if p.hdr_maximize else ' Maximize off' shared.log.debug(txt) # components = [{ k: getattr(v, 'device', None) } for k, v in model.components.items()] # shared.log.debug(f'Diffuser pipeline components: {components}') @@ -659,8 +664,15 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro # final decode since there is no refiner if not is_refiner_enabled: - if output is not None and output.images is not None and len(output.images) > 0: - results = vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality) + if output is not None: + if not hasattr(output, 'images') and hasattr(output, 'frames'): + shared.log.debug(f'Generated: frames={len(output.frames[0])}') + output.images = output.frames[0] + if output.images is not None and len(output.images) > 0: + results = vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality) + else: + shared.log.warning('Processing returned no results') + results = [] else: shared.log.warning('Processing returned no results') results = [] diff --git a/scripts/animatediff.py b/scripts/animatediff.py index 531335761..00100b3b8 100644 --- a/scripts/animatediff.py +++ b/scripts/animatediff.py @@ -128,7 +128,7 @@ class Script(scripts.Script): lora_index = gr.Dropdown(label='Lora', choices=list(LORAS), value='None') strength = gr.Slider(label='Strength', minimum=0.0, maximum=2.0, step=0.05, value=1.0) with gr.Row(): - latent_mode = gr.Checkbox(label='Latent mode', value=False) + latent_mode = gr.Checkbox(label='Latent mode', value=True, visible=False) with gr.Row(): video_type = gr.Dropdown(label='Video file', choices=['None', 'GIF', 'PNG', 'MP4'], value='None') duration = gr.Slider(label='Duration', minimum=0.25, maximum=10, step=0.25, value=2, visible=False)