Merge branch 'main' into JPPhoto-symmetry-enhancements

invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from contextlib import contextmanager
 from dataclasses import dataclass
 from math import ceil
@@ -6,6 +7,7 @@ from typing import Any, Callable, Dict, Optional, Union
 import numpy as np
 import torch
 from diffusers.models.cross_attention import AttnProcessor
+from einops import einops
 from typing_extensions import TypeAlias
 
 from invokeai.backend.globals import Globals
@@ -31,12 +33,17 @@ ModelForwardCallback: TypeAlias = Union[
 ]
 
 
+
+SymmetryType = Enum('SymmetryType', ['MIRROR', 'FADE'])
+
+
 @dataclass(frozen=True)
 class PostprocessingSettings:
     threshold: float
     warmup: float
     h_symmetry_time_pct: Optional[float]
     v_symmetry_time_pct: Optional[float]
+    symmetry_type: Optional[SymmetryType] = None
 
 
 class InvokeAIDiffuserComponent:
@@ -529,25 +536,44 @@ class InvokeAIDiffuserComponent:
         ):
             v_symmetry_time_pct = None
 
+        width = latents.shape[3]
+        height = latents.shape[2]
         dev = latents.device.type
+        dtype = latents.dtype
+        symmetry_type = postprocessing_settings.symmetry_type or SymmetryType.FADE
 
         latents.to(device="cpu")
 
+        def make_ramp(ease_in: int, total: int) -> torch.Tensor:
+            ramp1 = torch.linspace(start=1.0, end=0.5, steps=ease_in, device=dev)
+            ramp2 = torch.linspace(start=0.5, end=1.0, steps=total - ease_in, device=dev)
+            ramp = torch.cat((ramp1, ramp2))
+            return ramp
+
         if (
             h_symmetry_time_pct != None
             and self.last_percent_through < h_symmetry_time_pct
             and percent_through >= h_symmetry_time_pct
         ):
             # Horizontal symmetry occurs on the 3rd dimension of the latent
-            width = latents.shape[3]
             x_flipped = torch.flip(latents, dims=[3])
-            latents = torch.cat(
-                [
-                    latents[:, :, :, 0 : int(width / 2)],
-                    x_flipped[:, :, :, int(width / 2) : int(width)],
-                ],
-                dim=3,
-            )
+            if symmetry_type is SymmetryType.MIRROR:
+                latents = torch.cat(
+                    [
+                        latents[:, :, :, 0 : int(width / 2)],
+                        x_flipped[:, :, :, int(width / 2) : int(width)],
+                    ],
+                    dim=3,
+                )
+            elif symmetry_type is SymmetryType.FADE:
+                apply_width = width // 2
+                # Create a linear ramp so the middle gets perfect symmetry but the edges retain their original latents
+                ramp = make_ramp(ease_in=apply_width, total=width)
+                fade1 = einops.repeat(tensor=ramp, pattern='m -> 1 4 k m', k=height).to(latents.device).type(dtype)
+                fade0 = 1 - fade1
+                # Multiply the crossover region to retain details and avoid a "muddy" appearance
+                multiplier = (fade1 * fade0) * 1.25 + 1
+                latents = ((latents * fade1) + (x_flipped * fade0)) * multiplier
 
         if (
             v_symmetry_time_pct != None
@@ -555,15 +581,24 @@ class InvokeAIDiffuserComponent:
             and percent_through >= v_symmetry_time_pct
         ):
             # Vertical symmetry occurs on the 2nd dimension of the latent
-            height = latents.shape[2]
             y_flipped = torch.flip(latents, dims=[2])
-            latents = torch.cat(
-                [
-                    latents[:, :, 0 : int(height / 2)],
-                    y_flipped[:, :, int(height / 2) : int(height)],
-                ],
-                dim=2,
-            )
+            if symmetry_type is SymmetryType.MIRROR:
+                latents = torch.cat(
+                    [
+                        latents[:, :, 0 : int(height / 2)],
+                        y_flipped[:, :, int(height / 2) : int(height)],
+                    ],
+                    dim=2,
+                )
+            elif symmetry_type is SymmetryType.FADE:
+                apply_height = height // 2
+                # Create a linear ramp so the middle gets perfect symmetry but the edges retain their original latents
+                ramp = make_ramp(ease_in=apply_height, total=height)
+                fade1 = einops.repeat(tensor=ramp, pattern='m -> 1 4 m k', k=width).to(latents.device).type(dtype)
+                fade0 = 1 - fade1
+                # Multiply the crossover region to retain details and avoid a "muddy" appearance
+                multiplier = (fade1 * fade0) * 1.25 + 1
+                latents = ((latents * fade1) + (y_flipped * fade0)) * multiplier
 
         self.last_percent_through = percent_through
         return latents.to(device=dev)