animatediff full latent mode

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**

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:

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

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 = []

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)