hcpdiff 0.9.1__py3-none-any.whl → 2.1__py3-none-any.whl

hcpdiff/workflow/daam/hook.py

@@ -0,0 +1,109 @@
+from daam import AggregateHooker, RawHeatMapCollection, UNetCrossAttentionLocator, GlobalHeatMap
+from daam.trace import UNetCrossAttentionHooker
+from typing import List
+from diffusers import UNet2DConditionModel
+from PIL import Image
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+def auto_autocast(*args, **kwargs):
+    if not torch.cuda.is_available():
+        kwargs['enabled'] = False
+
+    return torch.cuda.amp.autocast(*args, **kwargs)
+
+class DiffusionHeatMapHooker(AggregateHooker):
+    def __init__(
+            self,
+            unet: UNet2DConditionModel,
+            tokenizer,
+            vae_scale_factor: int,
+            low_memory: bool = False,
+            load_heads: bool = False,
+            save_heads: bool = False,
+            data_dir: str = None
+    ):
+        self.all_heat_maps = RawHeatMapCollection()
+        h = (unet.config.sample_size * vae_scale_factor)
+        self.latent_hw = 4096 if h == 512 or h == 1024 else 9216  # 64x64 or 96x96 depending on if it's 2.0-v or 2.0
+        locate_middle = load_heads or save_heads
+        self.locator = UNetCrossAttentionLocator(restrict={0} if low_memory else None, locate_middle_block=locate_middle)
+        self.last_prompt: str = ''
+        self.last_image: Image.Image = None
+        self.time_idx = 0
+        self._gen_idx = 0
+
+        self.tokenizer = tokenizer
+
+        modules = [
+            UNetCrossAttentionHooker(
+                x,
+                self,
+                layer_idx=idx,
+                latent_hw=self.latent_hw,
+                load_heads=load_heads,
+                save_heads=save_heads,
+                data_dir=data_dir
+            ) for idx, x in enumerate(self.locator.locate(unet))
+        ]
+
+        super().__init__(modules)
+
+    def time_callback(self, *args, **kwargs):
+        self.time_idx += 1
+
+    @property
+    def layer_names(self):
+        return self.locator.layer_names
+
+    def compute_global_heat_map(self, prompt=None, factors=None, head_idxs: List[int]=None, layer_idx=None, normalize=False):
+        # type: (str, List[float], int, int, bool) -> GlobalHeatMap
+        """
+        Compute the global heat map for the given prompt, aggregating across time (inference steps) and space (different
+        spatial transformer block heat maps).
+
+        Args:
+            prompt: The prompt to compute the heat map for. If none, uses the last prompt that was used for generation.
+            factors: Restrict the application to heat maps with spatial factors in this set. If `None`, use all sizes.
+            head_idx: Restrict the application to heat maps with this head index. If `None`, use all heads.
+            layer_idx: Restrict the application to heat maps with this layer index. If `None`, use all layers.
+
+        Returns:
+            A heat map object for computing word-level heat maps.
+        """
+        heat_maps = self.all_heat_maps
+
+        if prompt is None:
+            prompt = self.last_prompt
+
+        if factors is None:
+            factors = {0, 1, 2, 4, 8, 16, 32, 64}
+        else:
+            factors = set(factors)
+
+        all_merges = []
+        x = int(np.sqrt(self.latent_hw))
+
+        with auto_autocast(dtype=torch.float32):
+            for (factor, layer, head), heat_map in heat_maps:
+                if (head_idxs is None or head in head_idxs) and (layer_idx is None or layer_idx == layer):
+                    heat_map = heat_map.unsqueeze(1)/25 # [L,1,H,W]
+                    # The clamping fixes undershoot.
+                    all_merges.append(F.interpolate(heat_map, size=(x, x), mode='bicubic').clamp_(min=0))
+
+            try:
+                maps = torch.stack(all_merges, dim=0) # [B*head, L, 1, H, W]
+            except RuntimeError:
+                if head_idxs is not None or layer_idx is not None:
+                    raise RuntimeError('No heat maps found for the given parameters.')
+                else:
+                    raise RuntimeError('No heat maps found. Did you forget to call `with trace(...)` during generation?')
+
+            maps = maps.mean(0)[:, 0] # [L,H,W]
+            #maps = maps[:len(self.tokenizer.tokenize(prompt)) + 2]  # 1 for SOS and 1 for padding
+
+            if normalize:
+                maps = maps / (maps[1:-1].sum(0, keepdim=True) + 1e-6)  # drop out [SOS] and [PAD] for proper probabilities
+
+        return GlobalHeatMap(self.tokenizer, prompt, maps)

hcpdiff/workflow/diffusion.py

@@ -1,209 +1,198 @@
-import inspect
+import random
+import warnings
 from typing import Dict, Any, Union, List
 
 import torch
+from hcpdiff.diffusion.sampler import BaseSampler, DiffusersSampler
+from hcpdiff.utils import prepare_seed
+from hcpdiff.utils.net_utils import get_dtype, to_cuda
+from rainbowneko.infer import BasicAction
 from torch.cuda.amp import autocast
 
-from .base import BasicAction, from_memory_context, MemoryMixin
-
 try:
     from diffusers.utils import randn_tensor
 except:
     # new version of diffusers
     from diffusers.utils.torch_utils import randn_tensor
 
-from hcpdiff.utils import prepare_seed
-from hcpdiff.utils.net_utils import get_dtype
-import random
-
 class InputFeederAction(BasicAction):
-    @from_memory_context
-    def __init__(self, ex_inputs: Dict[str, Any], unet=None):
-        super().__init__()
+    def __init__(self, ex_inputs: Dict[str, Any], key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
         self.ex_inputs = ex_inputs
-        self.unet = unet
 
-    def forward(self, **states):
-        if hasattr(self.unet, 'input_feeder'):
-            for feeder in self.unet.input_feeder:
-                feeder(self.ex_inputs)
-        return states
+    def forward(self, model, ex_inputs=None, **states):
+        ex_inputs = self.ex_inputs if ex_inputs is None else {**ex_inputs, **self.ex_inputs}
+        if hasattr(model, 'input_feeder'):
+            for feeder in model.input_feeder:
+                feeder(ex_inputs)
 
 class SeedAction(BasicAction):
-    @from_memory_context
-    def __init__(self, seed: Union[int, List[int]], bs: int = 1):
-        super().__init__()
+    def __init__(self, seed: Union[int, List[int]], bs: int = 1, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
         self.seed = seed
         self.bs = bs
 
-    def forward(self, device, **states):
+    def forward(self, device, gen_step=0, **states):
         bs = states['prompt_embeds'].shape[0]//2 if 'prompt_embeds' in states else self.bs
         if self.seed is None:
             seeds = [None]*bs
         elif isinstance(self.seed, int):
-            seeds = list(range(self.seed, self.seed+bs))
+            seeds = list(range(self.seed+gen_step*bs, self.seed+(gen_step+1)*bs))
         else:
             seeds = self.seed
         seeds = [s or random.randint(0, 1 << 30) for s in seeds]
 
         G = prepare_seed(seeds, device=device)
-        return {**states, 'seeds':seeds, 'generator':G, 'device':device}
+        return {'seeds':seeds, 'generator':G}
 
-class PrepareDiffusionAction(BasicAction, MemoryMixin):
-    def __init__(self, dtype='fp32', amp=True):
-        self.dtype = dtype
+class PrepareDiffusionAction(BasicAction):
+    def __init__(self, model_offload=False, amp=torch.float16, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
+        self.model_offload = model_offload
         self.amp = amp
 
-    def forward(self, memory, **states):
-        dtype = get_dtype(self.dtype)
-        memory.unet.to(dtype=dtype)
-        memory.text_encoder.to(dtype=dtype)
-        memory.vae.to(dtype=dtype)
+    def forward(self, device, denoiser, TE, vae, **states):
+        denoiser.to(device)
+        TE.to(device)
+        vae.to(device)
 
-        device = memory.unet.device
-        vae_scale_factor = 2**(len(memory.vae.config.block_out_channels)-1)
-        return {**states, 'dtype':self.dtype, 'amp':self.amp, 'device':device, 'vae_scale_factor':vae_scale_factor}
+        TE.eval()
+        denoiser.eval()
+        vae.eval()
+        return {'amp':self.amp, 'model_offload':self.model_offload}
 
-class MakeTimestepsAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, scheduler=None, N_steps: int = 30, strength: float = None):
-        self.scheduler = scheduler
+class MakeTimestepsAction(BasicAction):
+    def __init__(self, N_steps: int = 30, strength: float = None, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
         self.N_steps = N_steps
         self.strength = strength
 
-    def get_timesteps(self, timesteps, strength):
+    def get_timesteps(self, noise_sampler:BaseSampler, timesteps, strength):
         # get the original timestep using init_timestep
         num_inference_steps = len(timesteps)
         init_timestep = min(int(num_inference_steps*strength), num_inference_steps)
 
         t_start = max(num_inference_steps-init_timestep, 0)
-        timesteps = timesteps[t_start*self.scheduler.order:]
+        if isinstance(noise_sampler, DiffusersSampler):
+            timesteps = timesteps[t_start*noise_sampler.scheduler.order:]
+        else:
+            timesteps = timesteps[t_start:]
 
         return timesteps
 
-    def forward(self, memory, device, **states):
-        self.scheduler = self.scheduler or memory.scheduler
-
-        self.scheduler.set_timesteps(self.N_steps, device=device)
-        timesteps = self.scheduler.timesteps
+    def forward(self, noise_sampler:BaseSampler, device, **states):
+        timesteps = noise_sampler.get_timesteps(self.N_steps, device=device)
         if self.strength:
-            timesteps = self.get_timesteps(timesteps, self.strength)
-        alphas_cumprod = self.scheduler.alphas_cumprod.to(timesteps.device)
-        return {**states, 'device':device, 'timesteps':timesteps, 'alphas_cumprod':alphas_cumprod}
-
-class MakeLatentAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, scheduler=None, N_ch=4, height=512, width=512):
-        self.scheduler = scheduler
+            timesteps = self.get_timesteps(noise_sampler, timesteps, self.strength)
+            return {'timesteps':timesteps, 'start_timestep':timesteps[:1]}
+        else:
+            return {'timesteps':timesteps}
+
+class MakeLatentAction(BasicAction):
+    def __init__(self, N_ch=4, height=None, width=None, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
         self.N_ch = N_ch
         self.height = height
         self.width = width
 
-    def forward(self, memory, generator, device, dtype, bs=None, latents=None, vae_scale_factor=8, start_timestep=None, **states):
+    def forward(self, noise_sampler:BaseSampler, vae, generator, device, dtype, bs=None, latents=None, start_timestep=None,
+                pooled_output=None, crop_coord=None, **states):
         if bs is None:
             if 'prompt' in states:
                 bs = len(states['prompt'])
-        scheduler = self.scheduler or memory.scheduler
+        vae_scale_factor = 2**(len(vae.config.block_out_channels)-1)
+        device = torch.device(device)
 
-        shape = (bs, self.N_ch, self.height//vae_scale_factor, self.width//vae_scale_factor)
+        if latents is None:
+            shape = (bs, self.N_ch, self.height//vae_scale_factor, self.width//vae_scale_factor)
+        else:
+            if self.height is not None:
+                warnings.warn('latents exist! User-specified width and height will be ignored!')
+            shape = latents.shape
         if isinstance(generator, list) and len(generator) != bs:
             raise ValueError(
                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                 f" size of {bs}. Make sure the batch size matches the length of the generators."
             )
 
-        noise = randn_tensor(shape, generator=generator, device=device, dtype=get_dtype(dtype))
         if latents is None:
-            # scale the initial noise by the standard deviation required by the scheduler
-            latents = noise*scheduler.init_noise_sigma
+            # scale the initial noise by the standard deviation required by the noise_sampler
+            noise_sampler.generator = generator
+            latents = noise_sampler.init_noise(shape, device=device, dtype=get_dtype(dtype))
         else:
             # image to image
             latents = latents.to(device)
-            latents = scheduler.add_noise(latents, noise, start_timestep)
+            latents, noise = noise_sampler.add_noise(latents, start_timestep)
 
-        return {**states, 'latents':latents, 'device':device, 'dtype':dtype, 'generator':generator}
+        output = {'latents':latents}
 
-class NoisePredAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, unet=None, scheduler=None, guidance_scale: float = 7.0):
+        # SDXL inputs
+        if pooled_output is not None:
+            width, height = shape[3]*vae_scale_factor, shape[2]*vae_scale_factor
+            if crop_coord is None:
+                crop_info = torch.tensor([height, width, 0, 0, height, width], dtype=torch.float)
+            else:
+                crop_info = torch.tensor([height, width, *crop_coord], dtype=torch.float)
+            crop_info = crop_info.to(device).repeat(bs, 1)
+            output['text_embeds'] = pooled_output[-1].to(device)
+
+            if 'negative_prompt' in states:
+                output['crop_info'] = torch.cat([crop_info, crop_info], dim=0)
+
+        return output
+
+class DenoiseAction(BasicAction):
+    def __init__(self, guidance_scale: float = 7.0, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
         self.guidance_scale = guidance_scale
-        self.unet = unet
-        self.scheduler = scheduler
 
-    def forward(self, memory, t, latents, prompt_embeds, pooled_output=None, encoder_attention_mask=None, crop_info=None,
-                cross_attention_kwargs=None, dtype='fp32', amp=None, **states):
-        self.scheduler = self.scheduler or memory.scheduler
-        self.unet = self.unet or memory.unet
+    def forward(self, denoiser, noise_sampler: BaseSampler, t, latents, prompt_embeds, text_embeds=None, encoder_attention_mask=None, crop_info=None,
+                cross_attention_kwargs=None, dtype='fp32', amp=None, model_offload=False, **states):
+
+        if model_offload:
+            to_cuda(denoiser)  # to_cpu in VAE
 
         with autocast(enabled=amp is not None, dtype=get_dtype(amp)):
             latent_model_input = torch.cat([latents]*2) if self.guidance_scale>1 else latents
-            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            latent_model_input = noise_sampler.c_in(t)*latent_model_input
 
-            if pooled_output is None:
-                noise_pred = self.unet(latent_model_input, t, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
-                                       cross_attention_kwargs=cross_attention_kwargs, ).sample
+            if text_embeds is None:
+                noise_pred = denoiser(latent_model_input, t, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
+                                  cross_attention_kwargs=cross_attention_kwargs, ).sample
             else:
-                added_cond_kwargs = {"text_embeds":pooled_output, "time_ids":crop_info}
+                added_cond_kwargs = {"text_embeds":text_embeds, "time_ids":crop_info}
                 # predict the noise residual
-                noise_pred = self.unet(latent_model_input, t, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
-                                       cross_attention_kwargs=cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs).sample
+                noise_pred = denoiser(latent_model_input, t, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
+                                  cross_attention_kwargs=cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs).sample
 
             # perform guidance
             if self.guidance_scale>1:
                 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                 noise_pred = noise_pred_uncond+self.guidance_scale*(noise_pred_text-noise_pred_uncond)
 
-        return {**states, 'noise_pred':noise_pred, 'latents':latents, 't':t, 'prompt_embeds':prompt_embeds, 'pooled_output':pooled_output,
-            'crop_info':crop_info, 'cross_attention_kwargs':cross_attention_kwargs, 'dtype':dtype, 'amp':amp}
-
-class SampleAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, scheduler=None, eta=0.0):
-        self.scheduler = scheduler
-        self.eta = eta
-
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    def forward(self, memory, noise_pred, t, latents, generator, **states):
-        self.scheduler = self.scheduler or memory.scheduler
-
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, self.eta)
+        return {'noise_pred':noise_pred}
 
+class SampleAction(BasicAction):
+    def forward(self, noise_sampler: BaseSampler, noise_pred, t, latents, generator, **states):
         # compute the previous noisy sample x_t -> x_t-1
-        sc_out = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)
-        latents = sc_out.prev_sample
-        return {**states, 'latents':latents, 't':t, 'generator':generator}
-
-class DiffusionStepAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, unet=None, scheduler=None, guidance_scale: float = 7.0):
-        self.act_noise_pred = NoisePredAction(unet, scheduler, guidance_scale)
-        self.act_sample = SampleAction(scheduler)
-
-    def forward(self, memory, **states):
-        states = self.act_noise_pred(memory=memory, **states)
-        states = self.act_sample(memory=memory, **states)
+        latents = noise_sampler.denoise(latents, t, noise_pred, generator=generator)
+        return {'latents':latents}
+
+class DiffusionStepAction(BasicAction):
+    def __init__(self, guidance_scale: float = 7.0, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
+        self.act_noise_pred = DenoiseAction(guidance_scale)
+        self.act_sample = SampleAction()
+
+    def forward(self, denoiser, noise_sampler, **states):
+        states = self.act_noise_pred(denoiser=denoiser, noise_sampler=noise_sampler, **states)
+        states = self.act_sample(**states)
         return states
 
 class X0PredAction(BasicAction):
-    def forward(self, latents, alphas_cumprod, t, noise_pred, **states):
-        # x_t -> x_0
-        alpha_prod_t = alphas_cumprod[t.long()]
-        beta_prod_t = 1-alpha_prod_t
-        latents_x0 = (latents-beta_prod_t**(0.5)*noise_pred)/alpha_prod_t**(0.5)  # approximate x_0
-        return {**states, 'latents_x0':latents_x0, 'latents':latents, 'alphas_cumprod':alphas_cumprod, 't':t, 'noise_pred':noise_pred}
+    def forward(self, latents, noise_sampler: BaseSampler, t, noise_pred, **states):
+        latents_x0 = noise_sampler.eps_to_x0(noise_pred, latents, t)
+        return {'latents_x0':latents_x0}
+
+def time_iter(timesteps, **states):
+    return [{'t':t} for t in timesteps]

hcpdiff/workflow/fast.py

@@ -0,0 +1,31 @@
+from sfast.compilers.diffusion_pipeline_compiler import (compile_unet, CompilationConfig)
+from rainbowneko.infer import BasicAction
+
+
+class SFastCompileAction(BasicAction):
+
+    @staticmethod
+    def compile_model(unet):
+        # compile model
+        config = CompilationConfig.Default()
+        config.enable_xformers = False
+        try:
+            import xformers
+            config.enable_xformers = True
+        except ImportError:
+            print('xformers not installed, skip')
+        # NOTE:
+        # When GPU VRAM is insufficient or the architecture is too old, Triton might be slow.
+        # Disable Triton if you encounter this problem.
+        try:
+            import tritonx
+            config.enable_triton = True
+        except ImportError:
+            print('Triton not installed, skip')
+        config.enable_cuda_graph = True
+
+        return compile_unet(unet, config)
+
+    def forward(self, denoiser, **states):
+        denoiser = self.compile_model(denoiser)
+        return {'denoiser': denoiser}

hcpdiff/workflow/flow.py

@@ -0,0 +1,67 @@
+from rainbowneko.infer import BasicAction
+from typing import List, Dict
+from tqdm import tqdm
+import math
+
+class FilePromptAction(BasicAction):
+    def __init__(self, actions: List[BasicAction], prompt: str, negative_prompt: str, bs: int = 4, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
+        if prompt.endswith('.txt'):
+            with open(prompt, 'r') as f:
+                prompt = f.read().split('\n')
+        else:
+            prompt = [prompt]
+
+        if negative_prompt.endswith('.txt'):
+            with open(negative_prompt, 'r') as f:
+                negative_prompt = f.read().split('\n')
+        else:
+            negative_prompt = [negative_prompt]*len(prompt)
+
+        self.prompt = prompt
+        self.negative_prompt = negative_prompt
+        self.bs = bs
+        self.actions = actions
+
+
+    def forward(self, **states):
+        states.update({'prompt_all':self.prompt, 'negative_prompt_all':self.negative_prompt})
+        states_ref = dict(**states)
+
+        pbar = tqdm(range(math.ceil(len(self.prompt)/self.bs)))
+        N_steps = len(self.actions)
+        for gen_step in pbar:
+            states = dict(**states_ref)
+            feed_data = {'gen_step': gen_step}
+            states.update(feed_data)
+            for step, act in enumerate(self.actions):
+                pbar.set_description(f'[{step+1}/{N_steps}] action: {type(act).__name__}')
+                states = act(**states)
+        return states
+
+class FlowPromptAction(BasicAction):
+    def __init__(self, actions: List[BasicAction], prompt: str, negative_prompt: str, bs: int = 4, num: int = None, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
+        prompt = [prompt]*num
+        negative_prompt = [negative_prompt]*num
+
+        self.prompt = prompt
+        self.negative_prompt = negative_prompt
+        self.bs = bs
+        self.actions = actions
+
+
+    def forward(self, **states):
+        states.update({'prompt_all':self.prompt, 'negative_prompt_all':self.negative_prompt})
+        states_ref = dict(**states)
+
+        pbar = tqdm(range(math.ceil(len(self.prompt)/self.bs)))
+        N_steps = len(self.actions)
+        for gen_step in pbar:
+            states = dict(**states_ref)
+            feed_data = {'gen_step': gen_step}
+            states.update(feed_data)
+            for step, act in enumerate(self.actions):
+                pbar.set_description(f'[{step+1}/{N_steps}] action: {type(act).__name__}')
+                states = act(**states)
+        return states