hcpdiff 2.3.1__py3-none-any.whl → 2.4__py3-none-any.whl

hcpdiff/ckpt_manager/__init__.py

@@ -1,4 +1,4 @@
 from .format import EmbFormat, DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSDXLFormat, \
-    OfficialSD15Format, LoraWebuiFormat
+    OfficialSD15Format, LoraWebuiFormat, DiffusersFluxFormat, OneFileFluxFormat
 from .ckpt import EmbSaver
 from .loader import HCPLoraLoader

hcpdiff/ckpt_manager/format/__init__.py

@@ -1,4 +1,4 @@
 from .emb import EmbFormat
-from .diffusers import DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat
-from .sd_single import OfficialSD15Format, OfficialSDXLFormat
+from .diffusers import DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat, DiffusersFluxFormat
+from .sd_single import OfficialSD15Format, OfficialSDXLFormat, OneFileFluxFormat
 from .lora_webui import LoraWebuiFormat

hcpdiff/ckpt_manager/format/diffusers.py

@@ -1,10 +1,10 @@
 import torch
-from diffusers import ModelMixin, AutoencoderKL, UNet2DConditionModel, PixArtTransformer2DModel
+from diffusers import ModelMixin, AutoencoderKL, UNet2DConditionModel, PixArtTransformer2DModel, FluxTransformer2DModel
 from rainbowneko.ckpt_manager.format import CkptFormat
 from transformers import CLIPTextModel, AutoTokenizer, T5EncoderModel
 
-from hcpdiff.diffusion.sampler import VPSampler, DDPMDiscreteSigmaScheduler
-from hcpdiff.models.compose import SDXLTokenizer, SDXLTextEncoder
+from hcpdiff.diffusion.sampler import VPSampler, DDPMDiscreteSigmaScheduler, FlowSigmaScheduler, Sampler, FluxShiftTimeSampler
+from hcpdiff.models.compose import SDXLTokenizer, SDXLTextEncoder, FluxTokenizer, FluxTextEncoder
 
 class DiffusersModelFormat(CkptFormat):
     def __init__(self, builder: ModelMixin):
@@ -51,9 +51,24 @@ class DiffusersPixArtFormat(CkptFormat):
             pretrained_model, subfolder="transformer", revision=revision, torch_dtype=dtype
         )
         vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
-        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
+        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler(linear_start=0.0001, linear_end=0.02, beta_schedule='linear'))
 
         TE = TE or T5EncoderModel.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
         tokenizer = tokenizer or AutoTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
+        tokenizer.model_max_length = 300
 
         return dict(denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer)
+
+class DiffusersFluxFormat(CkptFormat):
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        denoiser = denoiser or FluxTransformer2DModel.from_pretrained(
+            pretrained_model, subfolder="transformer", revision=revision, torch_dtype=dtype
+        )
+        vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
+        noise_sampler = noise_sampler or Sampler(FlowSigmaScheduler(), t_sampler=FluxShiftTimeSampler())
+
+        TE = TE or FluxTextEncoder.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
+        tokenizer = tokenizer or FluxTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
+
+        return dict(denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer)

hcpdiff/ckpt_manager/format/emb.py

@@ -2,13 +2,18 @@ from typing import Tuple
 
 import torch
 from rainbowneko.ckpt_manager.format import CkptFormat
-from torch.serialization import FILE_LIKE
+from torch import nn, Tensor
+from rainbowneko.utils import FILE_LIKE
 
 class EmbFormat(CkptFormat):
     EXT = 'pt'
 
-    def save_ckpt(self, sd_model: Tuple[str, torch.Tensor], save_f: FILE_LIKE):
+    def save_ckpt(self, sd_model: Tuple[str, Tensor | nn.Parameter | nn.ParameterDict], save_f: FILE_LIKE):
         name, emb = sd_model
+        if hasattr(emb, 'named_parameters'):
+            emb = dict(emb.named_parameters())
+        elif isinstance(emb, nn.Parameter):
+            emb = emb.data
         torch.save({'string_to_param':{'*':emb}, 'name':name}, save_f)
 
     def load_ckpt(self, ckpt_f: FILE_LIKE, map_location="cpu"):
@@ -18,4 +23,4 @@ class EmbFormat(CkptFormat):
         else:
             emb = state['emb_params']
         emb.requires_grad_(False)
-        return emb
+        return emb

hcpdiff/ckpt_manager/format/lora_webui.py

@@ -3,7 +3,7 @@ import re
 from typing import List, Dict, Any
 
 from rainbowneko.ckpt_manager.format import CkptFormat, SafeTensorFormat
-from torch.serialization import FILE_LIKE
+from rainbowneko.utils import FILE_LIKE
 
 class LoraConverter:
     com_name_unet = ['down_blocks', 'up_blocks', 'mid_block', 'transformer_blocks', 'to_q', 'to_k', 'to_v', 'to_out', 'proj_in', 'proj_out',

hcpdiff/ckpt_manager/format/sd_single.py

@@ -1,9 +1,9 @@
 import torch
-from diffusers import AutoencoderKL, StableDiffusionPipeline, StableDiffusionXLPipeline
+from diffusers import AutoencoderKL, StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline
 from rainbowneko.ckpt_manager.format import CkptFormat
 
-from hcpdiff.diffusion.sampler import VPSampler, DDPMDiscreteSigmaScheduler
-from hcpdiff.models.compose import SDXLTextEncoder, SDXLTokenizer
+from hcpdiff.diffusion.sampler import VPSampler, DDPMDiscreteSigmaScheduler, FlowSigmaScheduler, Sampler, FluxShiftTimeSampler
+from hcpdiff.models.compose import SDXLTextEncoder, SDXLTokenizer, FluxTextEncoder, FluxTokenizer
 
 class OfficialSD15Format(CkptFormat):
     # Single file format
@@ -35,7 +35,30 @@ class OfficialSDXLFormat(CkptFormat):
         )
 
         noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
-        TE = SDXLTextEncoder([('clip_L', pipe.text_encoder), ('clip_bigG', pipe.text_encoder_2)])
-        tokenizer = SDXLTokenizer([('clip_L', pipe.tokenizer), ('clip_bigG', pipe.tokenizer_2)])
+        TE = SDXLTextEncoder({'clip_L': pipe.text_encoder, 'clip_bigG': pipe.text_encoder_2})
+        tokenizer = SDXLTokenizer({'clip_L': pipe.tokenizer, 'clip_bigG': pipe.tokenizer_2})
 
         return dict(denoiser=pipe.unet, TE=TE, vae=pipe.vae, noise_sampler=noise_sampler, tokenizer=tokenizer)
+
+class OneFileFluxFormat(CkptFormat):
+    # Single file format
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        pipe_args = dict(unet=denoiser, vae=vae)
+        if TE is not None:
+            pipe_args['text_encoder'] = TE.clip
+            pipe_args['text_encoder_2'] = TE.T5
+        if tokenizer is not None:
+            pipe_args['tokenizer'] = tokenizer.clip
+            pipe_args['tokenizer_2'] = tokenizer.T5
+
+        pipe_args = {k:v for k,v in pipe_args.items() if v is not None}
+        pipe = FluxPipeline.from_single_file(
+            pretrained_model, revision=revision, torch_dtype=dtype, **pipe_args
+        )
+
+        noise_sampler = noise_sampler or Sampler(FlowSigmaScheduler(), t_sampler=FluxShiftTimeSampler())
+        TE = FluxTextEncoder({'clip': pipe.text_encoder, 'T5': pipe.text_encoder_2})
+        tokenizer = FluxTokenizer({'clip': pipe.tokenizer, 'T5': pipe.tokenizer_2})
+
+        return dict(denoiser=pipe.unet, TE=TE, vae=pipe.vae, noise_sampler=noise_sampler, tokenizer=tokenizer)

hcpdiff/data/cache/vae.py

@@ -73,7 +73,11 @@ class VaeCache(DataCache):
                     for data in tqdm(loader):
                         image = data['image'].to(device=_share.device, dtype=vae.dtype)
                         latents = model.vae.encode(image).latent_dist.sample()
-                        latents = (latents*vae.config.scaling_factor).cpu()
+                        if shift_factor := getattr(vae.config, 'shift_factor', None) is not None:
+                            latents = (latents-shift_factor)*vae.config.scaling_factor
+                        else:
+                            latents = latents*vae.config.scaling_factor
+                        latents = latents.cpu()
 
                         for img_id, latent, coord in zip(data['id'], latents, data['coord']):
                             data_cache = {'latent': latent, 'coord': coord}
@@ -89,7 +93,11 @@ class VaeCache(DataCache):
                     img_id = data['id']
                     image = data['image'].to(device=_share.device, dtype=vae.dtype)
                     latents = model.vae.encode(image).latent_dist.sample()
-                    latents = (latents*vae.config.scaling_factor).cpu()
+                    if shift_factor := getattr(vae.config, 'shift_factor', None) is not None:
+                        latents = (latents-shift_factor)*vae.config.scaling_factor
+                    else:
+                        latents = latents*vae.config.scaling_factor
+                    latents = latents.cpu()
                     for img_id, latent, coord in zip(data['id'], latents, data['coord']):
                         self.cache[img_id] = {'latent': latent, 'coord': coord}
 

hcpdiff/data/handler/text.py

@@ -1,10 +1,12 @@
 import random
-from typing import Dict, Union, List
+from string import Formatter
+from typing import Dict, Union
 
 import numpy as np
-from string import Formatter
-from rainbowneko.data import DataHandler
 from rainbowneko._share import register_model_callback
+from rainbowneko.data import DataHandler
+
+from hcpdiff.models.compose import ComposeTokenizer
 
 class TagShuffleHandler(DataHandler):
     def __init__(self, key_map_in=('prompt -> prompt',), key_map_out=('prompt -> prompt',)):
@@ -58,7 +60,6 @@ class TagEraseHandler(DataHandler):
     def __repr__(self):
         return f'TagEraseHandler(p={self.p})'
 
-
 class TemplateFillHandler(DataHandler):
     def __init__(self, word_names: Dict[str, str], key_map_in=('prompt -> prompt',), key_map_out=('prompt -> prompt',)):
         super().__init__(key_map_in, key_map_out)
@@ -68,7 +69,7 @@ class TemplateFillHandler(DataHandler):
         template, caption = prompt['template'], prompt['caption']
 
         keys_need = {i[1] for i in Formatter().parse(template) if i[1] is not None}
-        fill_dict = {k: v for k, v in self.word_names.items() if k in keys_need}
+        fill_dict = {k:v for k, v in self.word_names.items() if k in keys_need}
 
         if (caption is not None) and ('caption' in keys_need):
             fill_dict.update(caption=fill_dict.get('caption', None) or caption)
@@ -96,16 +97,16 @@ class TokenizeHandler(DataHandler):
         self.tokenizer = model_wrapper.tokenizer
 
     def handle(self, prompt):
-        token_info = self.tokenizer(prompt, truncation=True, padding="max_length", return_tensors="pt",
-                                max_length=self.tokenizer.model_max_length*self.tokenizer.N_repeats)
-        tokens = token_info.input_ids.squeeze()
-        data = {'prompt':tokens}
-        if self.encoder_attention_mask and 'attention_mask' in token_info:
-            data['attn_mask'] = token_info.attention_mask.squeeze()
+        # Tokenizer: {'input_ids':Tensor, 'attention_mask':Tensor, 'position_ids':Tensor, ...}
+        # ComposeTokenizer: {'input_ids':{'model1':Tensor, 'model2':Tensor}, ...}
+        token_info = ComposeTokenizer.tokenize_ex(self.tokenizer, prompt, truncation=True, padding="max_length",
+                                                  return_tensors="pt", squeeze=True)
+        data = {'prompt':token_info.input_ids}
+        if 'attention_mask' in data:
+            data['attn_mask'] = data['attention_mask']
         if 'position_ids' in token_info:
-            data['position_ids'] = token_info.position_ids.squeeze()
-
+            data['position_ids'] = token_info['position_ids']
         return data
 
     def __repr__(self):
-        return f'TokenizeHandler(\nencoder_attention_mask={self.encoder_attention_mask}, tokenizer={self.tokenizer}\n)'
+        return f'TokenizeHandler(\nencoder_attention_mask={self.encoder_attention_mask}, tokenizer={self.tokenizer}\n)'

hcpdiff/diffusion/sampler/__init__.py

@@ -1,4 +1,5 @@
 from .sigma_scheduler import *
 from .base import BaseSampler, Sampler
 from .VP import VPSampler
-from .diffusers import DiffusersSampler
+from .diffusers import DiffusersSampler
+from .timer import *

hcpdiff/diffusion/sampler/base.py

@@ -1,8 +1,10 @@
 from typing import Tuple
 
 import torch
+from rainbowneko.utils import add_dims
 
 from .sigma_scheduler import SigmaScheduler
+from .timer import TimeSampler
 
 try:
     from diffusers.utils import randn_tensor
@@ -11,7 +13,8 @@ except:
     from diffusers.utils.torch_utils import randn_tensor
 
 class BaseSampler:
-    def __init__(self, sigma_scheduler: SigmaScheduler, pred_type='eps', target_type='eps', generator: torch.Generator = None):
+    def __init__(self, sigma_scheduler: SigmaScheduler, t_sampler:TimeSampler = None, pred_type='eps', target_type='eps',
+                 generator: torch.Generator = None):
         '''
         Some losses can only be calculated in a specific space. Such as SSIM in x0 space.
         The model pred need convert to target space.
@@ -19,6 +22,9 @@ class BaseSampler:
         :param pred_type: ['x0', 'eps', 'velocity', ..., None]  The output space of the model
         :param target_type: ['x0', 'eps', 'velocity', ..., None]  The space to calculate the loss
         '''
+        if t_sampler is None:
+            t_sampler = TimeSampler()
+        self.t_sampler = t_sampler
 
         self.sigma_scheduler = sigma_scheduler
         self.generator = generator
@@ -38,15 +44,20 @@ class BaseSampler:
 
     def add_noise(self, x, t) -> Tuple[torch.Tensor, torch.Tensor]:
         noise = self.make_nosie(x.shape, device=x.device)
-        alpha = self.sigma_scheduler.alpha(t).view(-1, 1, 1, 1).to(x.device)
-        sigma = self.sigma_scheduler.sigma(t).view(-1, 1, 1, 1).to(x.device)
+        alpha = add_dims(self.sigma_scheduler.alpha(t), x.ndim-1).to(x.device)
+        sigma = add_dims(self.sigma_scheduler.sigma(t), x.ndim-1).to(x.device)
         noisy_x = alpha*x+sigma*noise
         return noisy_x.to(dtype=x.dtype), noise.to(dtype=x.dtype)
 
-    def add_noise_rand_t(self, x):
-        bs = x.shape[0]
+    def add_noise_rand_t(self, x, reso=None):
+        if x.ndim == 3:
+            B,L,C = x.shape
+            reso = L if reso is None else reso
+        else:
+            B,C,H,W = x.shape
+            reso = H*W if reso is None else reso
         # timesteps: [0, 1]
-        timesteps = self.sigma_scheduler.sample(shape=(bs,))
+        timesteps = self.t_sampler.sample(shape=(B,), reso=reso)
         timesteps = timesteps.to(x.device)
         noisy_x, noise = self.add_noise(x, timesteps)
 

hcpdiff/diffusion/sampler/diffusers.py

@@ -22,8 +22,8 @@ class DiffusersSampler(BaseSampler):
 
     def c_in(self, t):
         one = torch.ones_like(t)
-        if hasattr(self.scheduler, '_step_index'):
-            self.scheduler._step_index = None
+        # if hasattr(self.scheduler, '_step_index'):
+        #     self.scheduler._step_index = None
         return self.scheduler.scale_model_input(one, t)
 
     def get_timesteps(self, N_steps, device='cuda'):
@@ -35,7 +35,8 @@ class DiffusersSampler(BaseSampler):
 
     def add_noise(self, x, t):
         noise = randn_tensor(x.shape, generator=self.generator, device=x.device, dtype=x.dtype)
-        return self.scheduler.add_noise(x, noise, t), noise
+        t_in = self.sigma_scheduler.c_noise(t)
+        return self.scheduler.add_noise(x, noise, t_in), 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

hcpdiff/diffusion/sampler/sigma_scheduler/base.py

@@ -7,21 +7,21 @@ class SigmaScheduler:
         return t
 
     def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
-        '''
+        r'''
         x(t) = \alpha(t)*x(0) + \sigma(t)*eps
         :param t: 0-1, rate of time step
         '''
         raise NotImplementedError
 
     def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
-        '''
+        r'''
         x(t) = \alpha(t)*x(0) + \sigma(t)*eps
         :param t: 0-1, rate of time step
         '''
         raise NotImplementedError
 
     def velocity(self, t: Union[float, torch.Tensor], dt=1e-8, normlize=True) -> Tuple[torch.Tensor, torch.Tensor]:
-        '''
+        r'''
         v(t) = dx(t)/dt = d\alpha(t)/dt * x(0) + d\sigma(t)/dt *eps
         :param t: 0-1, rate of time step
         :return: d\alpha(t)/dt, d\sigma(t)/dt
@@ -63,14 +63,14 @@ class SigmaScheduler:
             return torch.ones_like(t, dtype=torch.float32)
 
     def c_skip(self, t: Union[float, torch.Tensor]):
-        '''
+        r'''
         \hat{x}(0) = c_skip*x(t) + c_out*f(x(t))
         :param t: 0-1, rate of time step
         '''
         return 1./self.alpha(t)
 
     def c_out(self, t: Union[float, torch.Tensor]):
-        '''
+        r'''
         \hat{x}(0) = c_skip*x(t) + c_out*f(x(t))
         :param t: 0-1, rate of time step
         '''
@@ -78,12 +78,3 @@ class SigmaScheduler:
     
     def c_noise(self, t: Union[float, torch.Tensor]):
         return t
-
-    def sample(self, min_t=0.0, max_t=1.0, shape=(1,)) -> torch.Tensor:
-        if isinstance(min_t, float):
-            min_t = torch.full(shape, min_t)
-        if isinstance(max_t, float):
-            max_t = torch.full(shape, max_t)
-
-        t = torch.lerp(min_t, max_t, torch.rand_like(min_t))
-        return t

hcpdiff/diffusion/sampler/sigma_scheduler/ddpm.py

@@ -2,8 +2,9 @@ import math
 from typing import Union, Tuple, Callable
 
 import torch
-
 from hcpdiff.utils import invert_func
+from rainbowneko.utils import add_dims
+
 from .base import SigmaScheduler
 
 class DDPMDiscreteSigmaScheduler(SigmaScheduler):
@@ -93,15 +94,15 @@ class DDPMDiscreteSigmaScheduler(SigmaScheduler):
 
     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
+        return add_dims(self.posterior_mean_coef1[t].to(t.device), x_0.ndim-1)*x_0+add_dims(self.posterior_mean_coef2[t].to(t.device), x_t.ndim-1)*x_t
 
-    def get_post_log_var(self, t, x_t_var=None):
+    def get_post_log_var(self, t, ndim, 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)
+        min_log = add_dims(self.posterior_log_variance_clipped[t].to(t.device), ndim-1)
         if x_t_var is None:
             return min_log
         else:
-            max_log = self.betas.log()[t].view(-1, 1, 1, 1).to(t.device)
+            max_log = add_dims(self.betas.log()[t].to(t.device), ndim-1)
             # 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
@@ -201,7 +202,7 @@ class DDPMContinuousSigmaScheduler(SigmaScheduler):
         B = 1-beta_s
         B_At = B-A*t
 
-        # 避免数值不稳定
+        # eps for stable
         eps = 1e-12
         B = torch.clamp(B, min=eps)
         B_At = torch.clamp(B_At, min=eps)

hcpdiff/diffusion/sampler/sigma_scheduler/edm.py

@@ -21,7 +21,7 @@ class EDMSigmaScheduler(SigmaScheduler):
         return torch.lerp(min_inv_rho, max_inv_rho, t)**self.rho
 
     def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
-        '''
+        r'''
         x_t = c_in(t) * (x(0) + \sigma(t)*eps), eps~N(0,I)
         '''
         if isinstance(t, float):
@@ -31,7 +31,7 @@ class EDMSigmaScheduler(SigmaScheduler):
         return sigma_edm/torch.sqrt(sigma_edm**2+self.sigma_data**2)
 
     def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
-        '''
+        r'''
         x_t = c_in(t) * (x(0) + \sigma(t)*eps), eps~N(0,I)
         '''
         if isinstance(t, float):
@@ -41,7 +41,7 @@ class EDMSigmaScheduler(SigmaScheduler):
         return 1./torch.sqrt(sigma_edm**2+self.sigma_data**2)
 
     def c_skip(self, t: Union[float, torch.Tensor]):
-        '''
+        r'''
         \hat{x}(0) = c_skip(t)*(x(t)/c_in(t)) + c_out(t)*f(x(t))
         :param t: 0-1, rate of time step
         '''
@@ -49,7 +49,7 @@ class EDMSigmaScheduler(SigmaScheduler):
         return self.sigma_data**2/torch.sqrt(sigma_edm**2+self.sigma_data**2)
 
     def c_out(self, t: Union[float, torch.Tensor]):
-        '''
+        r'''
         \hat{x}(0) = c_skip(t)*(x(t)/c_in(t)) + c_out(t)*f(x(t))
         :param t: 0-1, rate of time step
         '''

hcpdiff/diffusion/sampler/sigma_scheduler/flow.py

@@ -23,7 +23,7 @@ class FlowSigmaScheduler(SigmaScheduler):
         return 1-t
 
     def velocity(self, t: Union[float, torch.Tensor], dt=1e-8, normlize=False) -> Tuple[torch.Tensor, torch.Tensor]:
-        '''
+        r'''
         v(t) = dx(t)/dt = d\alpha(t)/dt * x(0) + d\sigma(t)/dt *eps
         :param t: 0-1, rate of time step
         :return: d\alpha(t)/dt, d\sigma(t)/dt
@@ -59,14 +59,14 @@ class FlowSigmaScheduler(SigmaScheduler):
         return 1-sigma
 
     def c_skip(self, t: Union[float, torch.Tensor]):
-        '''
+        r'''
         \hat{x}(0) = c_skip*x(t) + c_out*f(x(t))
         :param t: 0-1, rate of time step
         '''
         return 1.
 
     def c_out(self, t: Union[float, torch.Tensor]):
-        '''
+        r'''
         \hat{x}(0) = c_skip*x(t) + c_out*f(x(t))
         :param t: 0-1, rate of time step
         '''

hcpdiff/diffusion/sampler/timer/__init__.py

@@ -0,0 +1,2 @@
+from .base import TimeSampler, LogitNormalSampler
+from .shift import ShiftTimeSampler, FluxShiftTimeSampler

hcpdiff/diffusion/sampler/timer/base.py

@@ -0,0 +1,26 @@
+import torch
+
+class TimeSampler:
+    def sample(self, min_t=0.0, max_t=1.0, shape=(1,), reso=0) -> torch.Tensor:
+        if isinstance(min_t, float):
+            min_t = torch.full(shape, min_t)
+        if isinstance(max_t, float):
+            max_t = torch.full(shape, max_t)
+
+        t = torch.lerp(min_t, max_t, torch.rand_like(min_t))
+        return t
+
+class LogitNormalSampler(TimeSampler):
+    def __init__(self, mean=0.0, std=1.0):
+        self.mean = mean
+        self.std = std
+
+    def sample(self, min_t=0.0, max_t=1.0, shape=(1,), reso=0) -> torch.Tensor:
+        if isinstance(min_t, float):
+            min_t = torch.full(shape, min_t)
+        if isinstance(max_t, float):
+            max_t = torch.full(shape, max_t)
+
+        t = torch.sigmoid(torch.normal(mean=self.mean, std=self.std, size=shape))
+        t = torch.lerp(min_t, max_t, t)
+        return t

hcpdiff/diffusion/sampler/timer/shift.py

@@ -0,0 +1,49 @@
+import torch
+import math
+from torch import Tensor
+
+from .base import TimeSampler
+
+class ShiftTimeSampler(TimeSampler):
+    def __init__(self, t_sampler: TimeSampler = None, base_reso=1024*1024):
+        self.t_sampler = t_sampler
+        self.base_reso = base_reso
+
+    def sample(self, min_t=0.0, max_t=1.0, shape=(1,), reso=0) -> torch.Tensor:
+        t = self.t_sampler.sample(min_t, max_t, shape)
+        shift = math.sqrt(self.base_reso/(reso))
+        t = (t*shift)/(1+(shift-1)*t)
+        return t
+
+class FluxShiftTimeSampler(TimeSampler):
+    def __init__(self, t_sampler: TimeSampler = None, base_shift: float = 0.5, max_shift: float = 1.15, base_reso=256, max_reso=4096):
+        self.t_sampler = t_sampler
+        self.base_shift = base_shift
+        self.max_shift = max_shift
+        self.base_reso = base_reso
+        self.max_reso = max_reso
+
+    def time_shift(self, mu: float|Tensor, sigma: float, t: Tensor):
+        if torch.is_tensor(mu):
+            mu = mu.to(t.device)
+            return torch.exp(mu)/(torch.exp(mu)+(1/t-1)**sigma)
+        else:
+            return math.exp(mu)/(math.exp(mu)+(1/t-1)**sigma)
+
+    def get_lin_function(self, xi, x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15):
+        '''
+        ^
+        |      .(x2,y2)
+        |     /
+        |   . (x1,y1)
+        |_________>
+        '''
+        m = (y2-y1)/(x2-x1)
+        b = y1-m*x1
+        return m*xi+b
+
+    def sample(self, min_t=0.0, max_t=1.0, shape=(1,), reso=0) -> torch.Tensor:
+        mu = self.get_lin_function(reso, x1=self.base_reso, y1=self.base_shift, x2=self.max_reso, y2=self.max_shift)
+        t = self.t_sampler.sample(min_t, max_t, shape)
+        t = self.time_shift(mu, 1.0, t)
+        return t

hcpdiff/easy/__init__.py

@@ -1,2 +1,3 @@
-from .model import SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader, ControlNet_SD15, make_controlnet_handler
+from .model import (SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader, ControlNet_SD15, make_controlnet_handler, Flux_auto_loader,
+                    auto_load_wrapper)
 from .sampler import Diffusers_SD

hcpdiff/easy/cfg/sd15_train.py

@@ -43,9 +43,7 @@ def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
             )
         ], weight_decay=1e-2),
 
-        ckpt_saver=dict(
-            SD15=ckpt_saver_dict
-        ),
+        ckpt_saver=ckpt_saver_dict,
 
         train=dict(
             train_steps=train_steps,

hcpdiff/easy/model/__init__.py

@@ -1,2 +1,2 @@
-from .loader import SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader
+from .loader import SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader, Flux_auto_loader, auto_load_wrapper
 from .cnet import ControlNet_SD15, make_controlnet_handler