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

hcpdiff/diffusion/sampler/ddpm.py

@@ -0,0 +1,20 @@
+import torch
+
+from .base import BaseSampler
+from .sigma_scheduler import SigmaScheduler
+
+class DDPMSampler(BaseSampler):
+    def __init__(self, sigma_scheduler: SigmaScheduler, generator: torch.Generator=None):
+        super().__init__(sigma_scheduler, generator)
+
+    def c_in(self, sigma):
+        return 1./(sigma**2+1).sqrt()
+
+    def c_out(self, sigma):
+        return -sigma
+
+    def c_skip(self, sigma):
+        return 1.
+
+    def denoise(self, x, sigma, eps=None, generator=None):
+        raise NotImplementedError

hcpdiff/diffusion/sampler/diffusers.py

@@ -0,0 +1,66 @@
+import torch
+import inspect
+from diffusers import SchedulerMixin, DDPMScheduler
+
+try:
+    from diffusers.utils import randn_tensor
+except:
+    # new version of diffusers
+    from diffusers.utils.torch_utils import randn_tensor
+
+from .base import BaseSampler
+from .sigma_scheduler import TimeSigmaScheduler
+
+class DiffusersSampler(BaseSampler):
+    def __init__(self, scheduler: SchedulerMixin, eta=0.0, generator: torch.Generator=None):
+        sigma_scheduler = TimeSigmaScheduler()
+        super().__init__(sigma_scheduler, generator)
+        self.scheduler = scheduler
+        self.eta = eta
+
+    def c_in(self, sigma):
+        one = torch.ones_like(sigma)
+        if hasattr(self.scheduler, '_step_index'):
+            self.scheduler._step_index = None
+        return self.scheduler.scale_model_input(one, sigma)
+
+    def c_out(self, sigma):
+        return -sigma
+
+    def c_skip(self, sigma):
+        if self.c_in(sigma) == 1.:  # DDPM model
+            return (sigma**2+1).sqrt()  # 1/sqrt(alpha_)
+        else:  # EDM model
+            return 1.
+
+    def get_timesteps(self, N_steps, device='cuda'):
+        self.scheduler.set_timesteps(N_steps, device=device)
+        return self.scheduler.timesteps
+
+    def init_noise(self, shape, device='cuda', dtype=torch.float32):
+        return randn_tensor(shape, generator=self.generator, device=device, dtype=dtype)*self.scheduler.init_noise_sigma
+
+    def add_noise(self, x, sigma):
+        noise = randn_tensor(x.shape, generator=self.generator, device=x.device, dtype=x.dtype)
+        return self.scheduler.add_noise(x, noise, sigma), noise
+
+    def prepare_extra_step_kwargs(self, scheduler, 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(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(scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def denoise(self, x_t, sigma, eps=None, generator=None):
+        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator, self.eta)
+        return self.scheduler.step(eps, sigma, x_t, **extra_step_kwargs).prev_sample

hcpdiff/diffusion/sampler/edm.py

@@ -0,0 +1,22 @@
+import torch
+
+from .base import BaseSampler
+from .sigma_scheduler import SigmaScheduler
+
+class EDMSampler(BaseSampler):
+    def __init__(self, sigma_scheduler: SigmaScheduler, generator: torch.Generator = None, sigma_data: float = 1.0, sigma_thr=1000):
+        super().__init__(sigma_scheduler, generator)
+        self.sigma_data = sigma_data
+        self.sigma_thr = sigma_thr
+
+    def c_in(self, sigma):
+        return 1/(sigma**2+self.sigma_data**2).sqrt()
+
+    def c_out(self, sigma):
+        return (sigma*self.sigma_data)/(sigma**2+self.sigma_data**2).sqrt()
+
+    def c_skip(self, sigma):
+        return self.sigma_data**2/(sigma**2+self.sigma_data**2)
+
+    def denoise(self, x, sigma, eps=None, generator=None):
+        raise NotImplementedError

hcpdiff/diffusion/sampler/sigma_scheduler/__init__.py

@@ -0,0 +1,3 @@
+from .base import SigmaScheduler
+from .ddpm import DDPMDiscreteSigmaScheduler, DDPMContinuousSigmaScheduler, TimeSigmaScheduler
+from .edm import EDMSigmaScheduler, EDMRefSigmaScheduler

hcpdiff/diffusion/sampler/sigma_scheduler/base.py

@@ -0,0 +1,14 @@
+from typing import Union, Tuple
+
+import torch
+
+class SigmaScheduler:
+
+    def get_sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        '''
+        :param t: 0-1, rate of time step
+        '''
+        raise NotImplementedError
+
+    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError

hcpdiff/diffusion/sampler/sigma_scheduler/ddpm.py

@@ -0,0 +1,197 @@
+import torch
+import math
+from typing import Union, Tuple
+from hcpdiff.utils import linear_interp
+from .base import SigmaScheduler
+
+class DDPMDiscreteSigmaScheduler(SigmaScheduler):
+    def __init__(self, beta_schedule: str = "scaled_linear", linear_start=0.00085, linear_end=0.0120, num_timesteps=1000):
+        super().__init__()
+        self.num_timesteps = num_timesteps
+        self.betas = self.make_betas(beta_schedule, linear_start, linear_end, num_timesteps)
+        alphas = 1.0-self.betas
+        self.alphas_cumprod = torch.cumprod(alphas, dim=0)
+        self.sigmas = ((1-self.alphas_cumprod)/self.alphas_cumprod).sqrt()
+
+        # for VLB calculation
+        self.alphas_cumprod_prev = torch.cat([alphas.new_tensor([1.0]), self.alphas_cumprod[:-1]])
+        self.posterior_mean_coef1 = self.betas*torch.sqrt(self.alphas_cumprod_prev)/(1.0-self.alphas_cumprod)
+        self.posterior_mean_coef2 = (1.0-self.alphas_cumprod_prev)*torch.sqrt(alphas)/(1.0-self.alphas_cumprod)
+
+        self.posterior_variance = self.betas*(1.0-self.alphas_cumprod_prev)/(1.0-self.alphas_cumprod)
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_log_variance_clipped = torch.log(torch.cat([self.posterior_variance[1:2], self.posterior_variance[1:]]))
+
+
+    @property
+    def sigma_min(self):
+        return self.sigmas[0]
+
+    @property
+    def sigma_max(self):
+        return self.sigmas[-1]
+
+    def get_sigma(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor(t)
+        return self.sigmas[(t*len(self.sigmas)).long()]
+
+    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
+        if isinstance(min_rate, float):
+            min_rate = torch.full(shape, min_rate)
+        if isinstance(max_rate, float):
+            max_rate = torch.full(shape, max_rate)
+
+        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
+        t_scale = (t*(self.num_timesteps-1e-5)).long()  # [0, num_timesteps-1)
+        return self.sigmas[t_scale], t
+
+    def sigma_to_t(self, sigma: Union[float, torch.Tensor]):
+        t = (self.sigmas-sigma).abs().argmin()
+        return t/self.num_timesteps
+
+    def get_post_mean(self, t, x_0, x_t):
+        t = (t*len(self.sigmas)).long()
+        return self.posterior_mean_coef1[t].view(-1, 1, 1, 1).to(t.device)*x_0 + self.posterior_mean_coef2[t].view(-1, 1, 1, 1).to(t.device)*x_t
+
+    def get_post_log_var(self, t, x_t_var=None):
+        t = (t*len(self.sigmas)).long()
+        min_log = self.posterior_log_variance_clipped[t].view(-1, 1, 1, 1).to(t.device)
+        if x_t_var is None:
+            return min_log
+        else:
+            max_log = self.betas.log()[t].view(-1, 1, 1, 1).to(t.device)
+            # The model_var_values is [-1, 1] for [min_var, max_var].
+            frac = (x_t_var+1)/2
+            model_log_variance = frac*max_log+(1-frac)*min_log
+            return model_log_variance
+
+
+    @staticmethod
+    def betas_for_alpha_bar(
+        num_diffusion_timesteps,
+        max_beta=0.999,
+        alpha_transform_type="cosine",
+    ):
+        """
+        Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+        (1-beta) over time from t = [0,1].
+
+        Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+        to that part of the diffusion process.
+
+
+        Args:
+            num_diffusion_timesteps (`int`): the number of betas to produce.
+            max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                         prevent singularities.
+            alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                         Choose from `cosine` or `exp`
+
+        Returns:
+            betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+        """
+        if alpha_transform_type == "cosine":
+
+            def alpha_bar_fn(t):
+                return math.cos((t+0.008)/1.008*math.pi/2)**2
+
+        elif alpha_transform_type == "exp":
+
+            def alpha_bar_fn(t):
+                return math.exp(t*-12.0)
+
+        else:
+            raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
+
+        betas = []
+        for i in range(num_diffusion_timesteps):
+            t1 = i/num_diffusion_timesteps
+            t2 = (i+1)/num_diffusion_timesteps
+            betas.append(min(1-alpha_bar_fn(t2)/alpha_bar_fn(t1), max_beta))
+        return torch.tensor(betas, dtype=torch.float32)
+
+    @staticmethod
+    def make_betas(beta_schedule, beta_start, beta_end, num_train_timesteps, betas=None):
+        if betas is not None:
+            return torch.tensor(betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            return torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            return torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32)**2
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            return DDPMDiscreteSigmaScheduler.betas_for_alpha_bar(num_train_timesteps)
+        elif beta_schedule == "sigmoid":
+            # GeoDiff sigmoid schedule
+            betas = torch.linspace(-6, 6, num_train_timesteps)
+            return torch.sigmoid(betas)*(beta_end-beta_start)+beta_start
+        else:
+            raise NotImplementedError(f"{beta_schedule} does is not implemented.")
+
+class DDPMContinuousSigmaScheduler(DDPMDiscreteSigmaScheduler):
+
+    def get_sigma(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor(t)
+        return linear_interp(self.sigmas, t)
+
+    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
+        if isinstance(min_rate, float):
+            min_rate = torch.full(shape, min_rate)
+        if isinstance(max_rate, float):
+            max_rate = torch.full(shape, max_rate)
+
+        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
+        t_scale = (t*(self.num_timesteps-1-1e-5))  # [0, num_timesteps-1)
+
+        return linear_interp(self.sigmas, t_scale), t
+
+    def sigma_to_t(self, sigma: Union[float, torch.Tensor]):
+        diff = self.sigmas-sigma
+        diff[diff<0] = float('inf')
+        t0 = diff.argmin().clamp(0, self.num_timesteps-2)
+        return t0 + diff.min()/(self.sigmas[t0+1]-self.sigmas[t0])
+
+class TimeSigmaScheduler(SigmaScheduler):
+    def __init__(self, num_timesteps=1000):
+        super().__init__()
+        self.num_timesteps = num_timesteps
+
+    def get_sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        '''
+        :param t: 0-1, rate of time step
+        '''
+        return t
+
+    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)) -> Tuple[torch.Tensor, torch.Tensor]:
+        if isinstance(min_rate, float):
+            min_rate = torch.full(shape, min_rate)
+        if isinstance(max_rate, float):
+            max_rate = torch.full(shape, max_rate)
+
+        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
+        t_scale = (t*(self.num_timesteps-1e-5)).long()  # [0, num_timesteps-1)
+        return t_scale, t
+
+if __name__ == '__main__':
+    from matplotlib import pyplot as plt
+    import numpy as np
+
+    sigma_scheduler = DDPMDiscreteSigmaScheduler()
+    print(sigma_scheduler.sigma_min, sigma_scheduler.sigma_max)
+    t = torch.linspace(0, 1, 1000)
+    rho = 1.
+    s2 = (sigma_scheduler.sigma_min**(1/rho)+t*(sigma_scheduler.sigma_max**(1/rho)-sigma_scheduler.sigma_min**(1/rho)))**rho
+    t2 = np.interp(s2.log().numpy(), sigma_scheduler.sigmas.log().numpy(), t.numpy())
+
+    plt.figure()
+    plt.plot(sigma_scheduler.sigmas)
+    plt.plot(t2*1000, s2)
+    plt.show()
+
+    plt.figure()
+    plt.plot(sigma_scheduler.sigmas.log())
+    plt.plot(t2*1000, s2.log())
+    plt.show()

hcpdiff/diffusion/sampler/sigma_scheduler/edm.py

@@ -0,0 +1,48 @@
+from typing import Union
+
+import torch
+import numpy as np
+
+from .base import SigmaScheduler
+
+class EDMSigmaScheduler(SigmaScheduler):
+    def __init__(self, sigma_min=0.002, sigma_max=80.0, rho=7.0, num_timesteps=1000):
+        self.sigma_min = torch.tensor(sigma_min)
+        self.sigma_max = torch.tensor(sigma_max)
+        self.rho = rho
+
+        self.num_timesteps=num_timesteps
+
+    def get_sigma(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor(t)
+
+        min_inv_rho = self.sigma_min**(1/self.rho)
+        max_inv_rho = self.sigma_max**(1/self.rho)
+        return torch.lerp(min_inv_rho, max_inv_rho, t)**self.rho
+
+    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
+        if isinstance(min_rate, float):
+            min_rate = torch.full(shape, min_rate)
+        if isinstance(max_rate, float):
+            max_rate = torch.full(shape, max_rate)
+
+        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
+        return self.get_sigma(t), t
+
+class EDMRefSigmaScheduler(EDMSigmaScheduler):
+    def __init__(self, ref_scheduler, sigma_min=0.002, sigma_max=80.0, rho=7.0, num_timesteps=1000):
+        super().__init__(sigma_min, sigma_max, rho, num_timesteps=num_timesteps)
+        self.ref_sigmas = ref_scheduler.sigmas.cpu().clip(min=1e-8).log().numpy()
+        self.ref_t = np.linspace(0, 1, len(self.ref_sigmas))
+
+    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
+        if isinstance(min_rate, float):
+            min_rate = torch.full(shape, min_rate)
+        if isinstance(max_rate, float):
+            max_rate = torch.full(shape, max_rate)
+
+        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
+        sigma = self.get_sigma(t)
+        t_rect = torch.tensor(np.interp(sigma.cpu().clip(min=1e-8).log().numpy(), self.ref_sigmas, self.ref_t))
+        return sigma, t_rect

hcpdiff/easy/__init__.py

@@ -0,0 +1,2 @@
+from .model import SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader, ControlNet_SD15, make_controlnet_handler
+from .sampler import Diffusers_SD

hcpdiff/easy/cfg/__init__.py

@@ -0,0 +1,3 @@
+from .sd15_train import SD15_lora_train, cfg_data_SD_ARB, cfg_data_SD_resize_crop, SD15_finetuning
+from .sdxl_train import SDXL_lora_train, SDXL_finetuning
+from .t2i import SD15_t2i, SDXL_t2i, SDXL_t2i_lora, SD15_t2i_lora, SDXL_t2i_parts, SD15_t2i_parts

hcpdiff/easy/cfg/sd15_train.py

@@ -0,0 +1,207 @@
+import torch
+from rainbowneko.ckpt_manager import ckpt_saver, LAYERS_TRAINABLE, NekoPluginSaver, SafeTensorFormat
+from rainbowneko.data import RatioBucket, FixedBucket
+from rainbowneko.parser import CfgWDPluginParser, neko_cfg, CfgWDModelParser, disable_neko_cfg
+from rainbowneko.utils import ConstantLR, Path_Like
+
+from hcpdiff.ckpt_manager import LoraWebuiFormat
+from hcpdiff.data import TextImagePairDataset, Text2ImageSource, StableDiffusionHandler
+from hcpdiff.data import VaeCache
+from hcpdiff.easy import SD15_auto_loader
+from hcpdiff.models import SD15Wrapper, TEHookCFG
+from hcpdiff.models.lora_layers_patch import LoraLayer
+
+@neko_cfg
+def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int = 500, lr: float = 1e-5, clip_skip: int = 0,
+                    dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0, name: str = 'SD15'):
+    if low_vram:
+        from bitsandbytes.optim import AdamW8bit
+        optimizer = AdamW8bit(_partial_=True)
+    else:
+        optimizer = torch.optim.AdamW(_partial_=True)
+
+    from cfgs.train.py import train_base, tuning_base
+
+    return dict(
+        _base_=[train_base, tuning_base],
+        mixed_precision=dtype,
+
+        model_part=CfgWDModelParser([
+            dict(
+                lr=lr,
+                layers=['denoiser'],  # train UNet
+            )
+        ], weight_decay=1e-2),
+
+        ckpt_saver=dict(
+            SD15=ckpt_saver(
+                ckpt_type='safetensors',
+                target_module='denoiser',
+                layers=LAYERS_TRAINABLE,
+            )
+        ),
+
+        train=dict(
+            train_steps=train_steps,
+            save_step=save_step,
+
+            optimizer=optimizer,
+
+            scheduler=ConstantLR(
+                _partial_=True,
+                warmup_steps=warmup_steps,
+            ),
+        ),
+
+        model=dict(
+            name=name,
+
+            ## Easy config
+            wrapper=SD15Wrapper.from_pretrained(
+                _partial_=True,
+                models=SD15_auto_loader(ckpt_path=base_model, _partial_=True),
+                TE_hook_cfg=TEHookCFG(clip_skip=clip_skip),
+            ),
+        ),
+
+        data_train=dataset,
+    )
+
+@neko_cfg
+def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int = 200, lr: float = 1e-4, rank: int = 4, alpha: float = None,
+                    clip_skip: int = 0, with_conv: bool = False, dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0,
+                    name: str = 'SD15', save_webui_format=False):
+    with disable_neko_cfg:
+        if alpha is None:
+            alpha = rank
+
+        if with_conv:
+            lora_layers = [
+                r're:denoiser.*\.attn.?$',
+                r're:denoiser.*\.ff$',
+                r're:denoiser.*\.resnets$',
+                r're:denoiser.*\.proj_in$',
+                r're:denoiser.*\.proj_out$',
+                r're:denoiser.*\.conv$',
+            ]
+        else:
+            lora_layers = [
+                r're:denoiser.*\.attn.?$',
+                r're:denoiser.*\.ff$',
+            ]
+
+    if low_vram:
+        from bitsandbytes.optim import AdamW8bit
+        optimizer = AdamW8bit(_partial_=True, betas=(0.9, 0.99))
+    else:
+        optimizer = torch.optim.AdamW(_partial_=True, betas=(0.9, 0.99))
+
+    if save_webui_format:
+        lora_format = LoraWebuiFormat()
+    else:
+        lora_format = SafeTensorFormat()
+
+    from cfgs.train.py.examples import SD_FT
+
+    return dict(
+        _base_=[SD_FT],
+        mixed_precision=dtype,
+
+        model_part=None,
+        model_plugin=CfgWDPluginParser(cfg_plugin=dict(
+            lora1=LoraLayer.wrap_model(
+                _partial_=True,
+                lr=lr,
+                rank=rank,
+                alpha=alpha,
+                layers=lora_layers
+            )
+        ), weight_decay=0.1),
+
+        ckpt_saver=dict(
+            _replace_ = True,
+            lora_unet=NekoPluginSaver(
+                format=lora_format,
+                target_plugin='lora1',
+            )
+        ),
+
+        train=dict(
+            train_steps=train_steps,
+            save_step=save_step,
+
+            optimizer=optimizer,
+
+            scheduler=ConstantLR(
+                _partial_=True,
+                warmup_steps=warmup_steps,
+            ),
+        ),
+
+        model=dict(
+            name=name,
+
+            wrapper=SD15Wrapper.from_pretrained(
+                _partial_=True,
+                models=SD15_auto_loader(ckpt_path=base_model, _partial_=True),
+                TE_hook_cfg=TEHookCFG(clip_skip=clip_skip),
+            ),
+        ),
+
+        data_train=dataset,
+    )
+
+@neko_cfg
+def cfg_data_SD_ARB(img_root: Path_Like, batch_size: int = 4, trigger_word: str = '', resolution: int = 512*512, num_bucket=4, word_names=None,
+                    prompt_dropout: float = 0, prompt_template: Path_Like = 'prompt_template/caption.txt', loss_weight=1.0):
+    if word_names is None:
+        word_names = dict(pt1=trigger_word)
+    else:
+        word_names = word_names
+
+    return TextImagePairDataset(
+        _partial_=True, batch_size=batch_size, loss_weight=loss_weight,
+        source=dict(
+            data_source1=Text2ImageSource(
+                img_root=img_root,
+                label_file='${.img_root}',  # path to image captions (file_words)
+                prompt_template=prompt_template,
+            ),
+        ),
+        handler=StableDiffusionHandler(
+            bucket=RatioBucket,
+            word_names=word_names,
+            erase=prompt_dropout,
+        ),
+        bucket=RatioBucket.from_files(
+            target_area=resolution,
+            num_bucket=num_bucket,
+        ),
+        cache=VaeCache(bs=batch_size)
+    )
+
+@neko_cfg
+def cfg_data_SD_resize_crop(img_root: Path_Like, batch_size: int = 4, trigger_word: str = '', target_size = (512, 512), word_names=None,
+                            prompt_dropout: float = 0, prompt_template: Path_Like = 'prompt_template/caption.txt', loss_weight=1.0):
+    if word_names is None:
+        word_names = dict(pt1=trigger_word)
+    else:
+        word_names = word_names
+
+    return TextImagePairDataset(
+        _partial_=True, batch_size=batch_size, loss_weight=loss_weight,
+        source=dict(
+            data_source1=Text2ImageSource(
+                img_root=img_root,
+                label_file='${.img_root}',  # path to image captions (file_words)
+                prompt_template=prompt_template,
+            ),
+        ),
+        handler=StableDiffusionHandler(
+            bucket=FixedBucket,
+            word_names=word_names,
+            erase=prompt_dropout,
+        ),
+        bucket=FixedBucket(target_size=target_size),
+        cache=VaeCache(bs=batch_size)
+    )