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

hcpdiff/workflow/diffusion.py

@@ -1,8 +1,13 @@
-from .base import BasicAction, from_memory_context, ExecAction, MemoryMixin
+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
-import inspect
 
 try:
     from diffusers.utils import randn_tensor
@@ -10,197 +15,184 @@ 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, to_cpu, to_cuda
-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}
-
-class PrepareDiffusionAction(BasicAction, MemoryMixin):
-    def __init__(self, dtype='fp32'):
-        self.dtype = dtype
-
-    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)
-
-        device = memory.unet.device
-        vae_scale_factor = 2**(len(memory.vae.config.block_out_channels)-1)
-        return {**states, 'dtype': self.dtype, 'device':device, 'vae_scale_factor':vae_scale_factor}
-
-class MakeTimestepsAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, scheduler=None, N_steps:int=30, strength:float=None):
-        self.scheduler = scheduler
+        return {'seeds':seeds, 'generator':G}
+
+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, device, denoiser, TE, vae, **states):
+        denoiser.to(device)
+        TE.to(device)
+        vae.to(device)
+
+        TE.eval()
+        denoiser.eval()
+        vae.eval()
+        return {'amp':self.amp, 'model_offload':self.model_offload}
+
+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)
+        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 :]
+        t_start = max(num_inference_steps-init_timestep, 0)
+        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
-        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):
+            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, 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)
+
+        output = {'latents':latents}
+
+        # 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)
 
-        return {**states, 'latents': latents, 'device':device, 'dtype':dtype, 'generator':generator}
+            if 'negative_prompt' in states:
+                output['crop_info'] = torch.cat([crop_info, crop_info], dim=0)
 
-class NoisePredAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, unet=None, scheduler=None, guidance_scale:float=7.0):
-        self.guidance_scale=guidance_scale
-        self.unet = unet
-        self.scheduler = scheduler
+        return output
 
-    def forward(self, memory, t, latents, prompt_embeds, pooled_output=None, encoder_attention_mask=None, crop_info=None,
-                cross_attention_kwargs=None, dtype='fp32', **states):
-        self.scheduler = self.scheduler or memory.scheduler
-        self.unet = self.unet or memory.unet
+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
 
-        with autocast(enabled=dtype == 'amp'):
+    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}
-
-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

hcpdiff/workflow/io.py

@@ -1,88 +1,56 @@
 import os
+from functools import partial
+from typing import List, Union
 
-from diffusers import UNet2DConditionModel, AutoencoderKL, PNDMScheduler
-
-from hcpdiff.utils import auto_text_encoder, auto_tokenizer, to_validate_file
-from hcpdiff.utils.cfg_net_tools import HCPModelLoader
-from hcpdiff.utils.img_size_tool import types_support
+import torch
+from hcpdiff.utils import to_validate_file
 from hcpdiff.utils.net_utils import get_dtype
-from .base import BasicAction, from_memory_context, MemoryMixin
-
-class LoadModelsAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, pretrained_model: str, dtype: str, unet=None, text_encoder=None, tokenizer=None, vae=None, scheduler=None):
-        self.pretrained_model = pretrained_model
+from rainbowneko.ckpt_manager import NekoLoader
+from rainbowneko.infer import BasicAction
+from rainbowneko.infer import LoadImageAction as Neko_LoadImageAction
+from rainbowneko.utils.img_size_tool import types_support
+
+class BuildModelsAction(BasicAction):
+    def __init__(self, model_loader: partial[NekoLoader.load], dtype: str=torch.float32, device='cuda', key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
+        self.model_loader = model_loader
         self.dtype = get_dtype(dtype)
+        self.device = device
 
-        self.unet = unet
-        self.text_encoder = text_encoder
-        self.tokenizer = tokenizer
-        self.vae = vae
-        self.scheduler = scheduler
+    def forward(self, in_preview=False, model=None, **states):
+        if in_preview:
+            model = self.model_loader(dtype=self.dtype, device=self.device, denoiser=model.denoiser, TE=model.TE, vae=model.vae)
+        else:
+            model = self.model_loader(dtype=self.dtype, device=self.device)
 
-    def forward(self, memory, **states):
-        memory.unet = self.unet or UNet2DConditionModel.from_pretrained(self.pretrained_model, subfolder="unet", torch_dtype=self.dtype)
-        memory.text_encoder = self.text_encoder or auto_text_encoder(self.pretrained_model, subfolder="text_encoder", torch_dtype=self.dtype)
-        memory.tokenizer = self.tokenizer or auto_tokenizer(self.pretrained_model, subfolder="tokenizer", use_fast=False)
-        memory.vae = self.vae or AutoencoderKL.from_pretrained(self.pretrained_model, subfolder="vae", torch_dtype=self.dtype)
-        memory.scheduler = self.scheduler or PNDMScheduler.from_pretrained(self.pretrained_model, subfolder="scheduler", torch_dtype=self.dtype)
+        if isinstance(model, dict):
+            return model
+        else:
+            return {'model':model}
 
-        return states
+class LoadImageAction(Neko_LoadImageAction):
+    def __init__(self, image_paths: Union[str, List[str]], image_transforms=None, key_map_in=None, key_map_out=('input.x -> images',)):
+        super().__init__(image_paths, image_transforms, key_map_in, key_map_out)
 
 class SaveImageAction(BasicAction):
-    @from_memory_context
-    def __init__(self, save_root: str, image_type: str = 'png', quality: int = 95):
+    def __init__(self, save_root: str, image_type: str = 'png', quality: int = 95, save_cfg=True, key_map_in=None, key_map_out=None):
+        super().__init__(key_map_in, key_map_out)
         self.save_root = save_root
         self.image_type = image_type
         self.quality = quality
+        self.save_cfg = save_cfg
 
         os.makedirs(save_root, exist_ok=True)
 
-    def forward(self, images, prompt, negative_prompt, seeds=None, **states):
-        num_img_exist = max([0]+[int(x.split('-', 1)[0]) for x in os.listdir(self.save_root) if x.rsplit('.', 1)[-1] in types_support])+1
+    def forward(self, images, prompt, negative_prompt, seeds, cfgs=None, parser=None, preview_root=None, preview_step=None, **states):
+        save_root = preview_root or self.save_root
+        num_img_exist = max([0]+[int(x.split('-', 1)[0]) for x in os.listdir(save_root) if x.rsplit('.', 1)[-1] in types_support])+1
 
         for bid, (p, pn, img) in enumerate(zip(prompt, negative_prompt, images)):
-            img_path = os.path.join(self.save_root, f"{num_img_exist}-{seeds[bid]}-{to_validate_file(prompt[0])}.{self.image_type}")
+            img_path = os.path.join(save_root, f"{preview_step or num_img_exist}-{seeds[bid]}-{to_validate_file(prompt[0])}.{self.image_type}")
             img.save(img_path, quality=self.quality)
             num_img_exist += 1
 
-        return {**states, 'images':images, 'prompt':prompt, 'negative_prompt':negative_prompt, 'seeds':seeds}
-
-class BuildModelLoaderAction(BasicAction, MemoryMixin):
-    def forward(self, memory, **states):
-        memory.model_loader_unet = HCPModelLoader(memory.unet)
-        memory.model_loader_TE = HCPModelLoader(memory.text_encoder)
-        return states
-
-class LoadPartAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, model: str, cfg):
-        self.model = model
-        self.cfg = cfg
-
-    def forward(self, memory, **states):
-        model_loader = memory[f"model_loader_{self.model}"]
-        model_loader.load_part(self.cfg)
-        return states
-
-class LoadLoraAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, model: str, cfg):
-        self.model = model
-        self.cfg = cfg
-
-    def forward(self, memory, **states):
-        model_loader = memory[f"model_loader_{self.model}"]
-        model_loader.load_lora(self.cfg)
-        return states
-
-class LoadPluginAction(BasicAction, MemoryMixin):
-    @from_memory_context
-    def __init__(self, model: str, cfg):
-        self.model = model
-        self.cfg = cfg
-
-    def forward(self, memory, **states):
-        model_loader = memory[f"model_loader_{self.model}"]
-        model_loader.load_plugin(self.cfg)
-        return states
+            if self.save_cfg:
+                cfgs.seed = seeds[bid]
+                parser.save_configs(cfgs, os.path.join(save_root, f"{preview_step or num_img_exist}-{seeds[bid]}-info"))