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

hcpdiff/tools/dataset_generator.py

@@ -0,0 +1,94 @@
+import argparse
+import json
+import os.path
+from typing import Callable
+
+import pyarrow.parquet as pq
+import torch
+from PIL import Image
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from tqdm.auto import tqdm
+
+from hcpdiff.data.caption_loader import auto_caption_loader
+
+class DatasetCreator:
+    def __init__(self, pretrained_model, out_dir: str, img_w: int=512, img_h: int=512):
+        scheduler = DPMSolverMultistepScheduler(
+            beta_start = 0.00085,
+            beta_end = 0.012,
+            beta_schedule = 'scaled_linear',
+            algorithm_type = 'dpmsolver++',
+            use_karras_sigmas = True,
+        )
+
+        self.pipeline = DiffusionPipeline.from_pretrained(pretrained_model, scheduler=scheduler, torch_dtype=torch.float16)
+        self.pipeline.requires_safety_checker = False
+        self.pipeline.safety_checker = None
+        self.pipeline.to("cuda")
+        self.pipeline.unet.to(memory_format=torch.channels_last)
+        #self.pipeline.enable_xformers_memory_efficient_attention()
+
+        self.out_dir = out_dir
+        self.img_w = img_w
+        self.img_h = img_h
+
+    def create_from_prompt_dataset(self, prompt_file: str, negative_prompt: str, bs: int, num: int=None, repeat:int=1, save_fmt:str='txt',
+                                   callback: Callable[[int, int], bool] = None):
+        os.makedirs(self.out_dir, exist_ok=True)
+        data = auto_caption_loader(prompt_file).load()
+        data = list(data.items())
+        data = self.split_batch(data, bs)  # [[(k,v),...],...]
+
+        if num is None:
+            num = len(data)
+        total = num*bs
+        count = 0
+        captions = {}
+        with torch.inference_mode():
+            for i in tqdm(range(num)):
+                for r in range(repeat):
+                    name_batch, p_batch = list(zip(*data[i%len(data)]))
+                    imgs = self.pipeline(list(p_batch), negative_prompt=[negative_prompt]*len(p_batch), num_inference_steps=25,
+                                         width=self.img_w, height=self.img_h).images
+                    for name, prompt, img in zip(name_batch, p_batch, imgs):
+                        img.save(os.path.join(self.out_dir, f'{count}_{name}.png'), format='PNG')
+                        captions[f'{count}_{name}'] = prompt
+                        count += 1
+                    if callback:
+                        if not callback(count, total):
+                            break
+
+        if save_fmt=='txt':
+            for k, v in captions.items():
+                with open(os.path.join(self.out_dir, f'{k}.txt'), "w") as f:
+                    f.write(v)
+        elif save_fmt=='json':
+            with open(os.path.join(self.out_dir, f'image_captions.json'), "w") as f:
+                json.dump(captions, f)
+        else:
+            raise ValueError(f"Invalid save_fmt: {save_fmt}")
+
+    @staticmethod
+    def split_batch(data, bs):
+        return [data[i:i+bs] for i in range(0, len(data), bs)]
+
+# python dataset_generator.py --prompt_file 标注文件或文件夹路径 --model 模型名称 --out_dir 输出文件夹路径 --repeat 每个prompt生成几个图 --bs batch_size --img_w 图片宽度 --img_h 图片高度
+# python dataset_generator.py --prompt_file 标注文件或文件夹路径 --model 模型名称 --out_dir 输出文件夹路径 --repeat 1 --bs 4 --img_w 640 --img_h 640
+if __name__ == '__main__':
+    torch.backends.cudnn.benchmark = True
+    parser = argparse.ArgumentParser(description='Diffusion Dataset Generator')
+    parser.add_argument('--prompt_file', type=str, default='')
+    parser.add_argument('--model', type=str, default='runwayml/stable-diffusion-v1-5')
+    parser.add_argument('--out_dir', type=str, default=r'./prompt_ds')
+    parser.add_argument('--negative_prompt', type=str,
+                        default='lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry')
+    parser.add_argument('--num', type=int, default=200)
+    parser.add_argument('--repeat', type=int, default=1)
+    parser.add_argument('--save_fmt', type=str, default='txt')
+    parser.add_argument('--bs', type=int, default=4)
+    parser.add_argument('--img_w', type=int, default=512)
+    parser.add_argument('--img_h', type=int, default=512)
+    args = parser.parse_args()
+
+    ds_creator = DatasetCreator(args.model, args.out_dir, args.img_w, args.img_h)
+    ds_creator.create_from_prompt_dataset(args.prompt_file, args.negative_prompt, args.bs, args.num, repeat=args.repeat, save_fmt=args.save_fmt)

hcpdiff/tools/download_hf_model.py

@@ -0,0 +1,24 @@
+from diffusers import DiffusionPipeline
+import argparse
+import torch
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Download Model')
+    parser.add_argument('--model', type=str, default='runwayml/stable-diffusion-v1-5')
+    parser.add_argument("--fp16", default=False, action="store_true")
+    parser.add_argument("--use_safetensors", default=False, action="store_true")
+    parser.add_argument("--out_path", type=str, default='ckpts/sd15')
+    args = parser.parse_args()
+
+    load_args = dict(torch_dtype = torch.float16 if args.fp16 else torch.float32)
+    save_args = dict()
+
+    if args.fp16:
+        load_args['variant'] = "fp16"
+        save_args['variant'] = "fp16"
+    if args.use_safetensors:
+        load_args['use_safetensors'] = True
+        save_args['safe_serialization'] = True
+
+    pipe = DiffusionPipeline.from_pretrained(args.model, **load_args)
+    pipe.save_pretrained(args.out_path, **save_args)

hcpdiff/tools/init_proj.py

@@ -1,23 +1,5 @@
-import sys
-import shutil
-import os
+from rainbowneko.tools.init_proj import copy_package_data
 
 def main():
-    prefix = sys.prefix
-    if not os.path.exists(os.path.join(prefix, 'hcpdiff')):
-        prefix = os.path.join(prefix, 'local')
-    try:
-        if os.path.exists(r'./cfgs'):
-            shutil.rmtree(r'./cfgs')
-        if os.path.exists(r'./prompt_tuning_template'):
-            shutil.rmtree(r'./prompt_tuning_template')
-        shutil.copytree(os.path.join(prefix, 'hcpdiff/cfgs'), r'./cfgs')
-        shutil.copytree(os.path.join(prefix, 'hcpdiff/prompt_tuning_template'), r'./prompt_tuning_template')
-    except:
-        try:
-            shutil.copytree(os.path.join(prefix, '../hcpdiff/cfgs'), r'./cfgs')
-            shutil.copytree(os.path.join(prefix, '../hcpdiff/prompt_tuning_template'), r'./prompt_tuning_template')
-        except:
-            this_file_dir = os.path.dirname(os.path.abspath(__file__))
-            shutil.copytree(os.path.join(this_file_dir, '../../cfgs'), r'./cfgs')
-            shutil.copytree(os.path.join(this_file_dir, '../../prompt_tuning_template'), r'./prompt_tuning_template')
+    copy_package_data('hcpdiff', 'cfgs', './cfgs')
+    copy_package_data('hcpdiff', 'prompt_template', './prompt_template')

hcpdiff/tools/lora_convert.py

@@ -3,15 +3,14 @@ import os.path
 from typing import List
 import math
 
-from hcpdiff.ckpt_manager import auto_manager
-from hcpdiff.deprecated import convert_to_webui_maybe_old, convert_to_webui_xl_maybe_old
+from rainbowneko.ckpt_manager import auto_ckpt_loader, NekoModelSaver
 
 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', 'input_blocks', 'middle_block', 'output_blocks']
     com_name_TE = ['self_attn', 'q_proj', 'v_proj', 'k_proj', 'out_proj', 'text_model']
     prefix_unet = 'lora_unet_'
     prefix_TE = 'lora_te_'
-    prefix_TE_xl_clip_B = 'lora_te1_'
+    prefix_TE_xl_clip_L = 'lora_te1_'
     prefix_TE_xl_clip_bigG = 'lora_te2_'
 
     lora_w_map = {'lora_down.weight': 'W_down', 'lora_up.weight':'W_up'}
@@ -26,14 +25,14 @@ class LoraConverter:
             sd_TE = self.convert_from_webui_(state, prefix=self.prefix_TE, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
         else:
             sd_unet = self.convert_from_webui_xl_unet_(state, prefix=self.prefix_unet, com_name=self.com_name_unet, com_name_tmp=self.com_name_unet_tmp)
-            sd_TE = self.convert_from_webui_xl_te_(state, prefix=self.prefix_TE_xl_clip_B, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
+            sd_TE = self.convert_from_webui_xl_te_(state, prefix=self.prefix_TE_xl_clip_L, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
             sd_TE2 = self.convert_from_webui_xl_te_(state, prefix=self.prefix_TE_xl_clip_bigG, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
             sd_TE.update(sd_TE2)
 
         if auto_scale_alpha:
             sd_unet = self.alpha_scale_from_webui(sd_unet)
             sd_TE = self.alpha_scale_from_webui(sd_TE)
-        return {'lora': sd_TE},  {'lora': sd_unet}
+        return {'plugin': sd_TE},  {'plugin': sd_unet}
 
     def convert_to_webui(self, sd_unet, sd_TE, auto_scale_alpha=False, sdxl=False):
         sd_unet = self.convert_to_webui_(sd_unet, prefix=self.prefix_unet)
@@ -59,7 +58,6 @@ class LoraConverter:
                 sd_covert[f'{model_k}.___.layer.{self.lora_w_map[lora_k]}'] = v
         return sd_covert
 
-    @convert_to_webui_maybe_old
     def convert_to_webui_(self, state, prefix):
         sd_covert = {}
         for k, v in state.items():
@@ -75,7 +73,6 @@ class LoraConverter:
             sd_covert[f"{prefix}{model_k.replace('.', '_')}.{lora_k}"] = v
         return sd_covert
 
-    @convert_to_webui_xl_maybe_old
     def convert_to_webui_xl_(self, state, prefix):
         sd_convert = {}
         for k, v in state.items():
@@ -90,7 +87,7 @@ class LoraConverter:
 
             new_k = f"{prefix}{model_k.replace('.', '_')}.{lora_k}"
             if 'clip' in new_k:
-                new_k = new_k.replace('_clip_B', '1') if 'clip_B' in new_k else new_k.replace('_clip_bigG', '2')
+                new_k = new_k.replace('_clip_L', '1') if 'clip_L' in new_k else new_k.replace('_clip_bigG', '2')
             sd_convert[new_k] = v
         return sd_convert
     
@@ -103,7 +100,7 @@ class LoraConverter:
             model_k, lora_k = k[prefix_len:].split('.', 1)
             model_k = self.replace_all(model_k, com_name, com_name_tmp).replace('_', '.').replace('%', '_')
             if prefix == 'lora_te1_':
-                model_k = f'clip_B.{model_k}'
+                model_k = f'clip_L.{model_k}'
             else:
                 model_k = f'clip_bigG.{model_k}'
 
@@ -224,23 +221,27 @@ if __name__ == '__main__':
 
     # load lora model
     print('convert lora model')
-    ckpt_manager = auto_manager(args.lora_path)
+    ckpt_loader = auto_ckpt_loader(args.lora_path)
+    ckpt_saver = NekoModelSaver(
+        format=ckpt_loader.format,
+        source=ckpt_loader.source,
+    )
 
     if args.from_webui:
-        state = ckpt_manager.load_ckpt(args.lora_path)
+        state = ckpt_loader.load(args.lora_path)
         # convert the weight name
         sd_TE, sd_unet = converter.convert_from_webui(state, auto_scale_alpha=args.auto_scale_alpha, sdxl=args.sdxl)
         # wegiht save
         os.makedirs(args.dump_path, exist_ok=True)
         TE_path = os.path.join(args.dump_path, 'TE-'+lora_name)
         unet_path = os.path.join(args.dump_path, 'unet-'+lora_name)
-        ckpt_manager._save_ckpt(sd_TE, save_path=TE_path)
-        ckpt_manager._save_ckpt(sd_unet, save_path=unet_path)
+        ckpt_saver.save(sd_TE, TE_path)
+        ckpt_saver.save(sd_unet, unet_path)
         print('save text encoder lora to:', TE_path)
         print('save unet lora to:', unet_path)
     elif args.to_webui:
-        sd_unet = ckpt_manager.load_ckpt(args.lora_path)
-        sd_TE = ckpt_manager.load_ckpt(args.lora_path_TE) if args.lora_path_TE else {'lora':{}}
-        state = converter.convert_to_webui(sd_unet['lora'], sd_TE['lora'], auto_scale_alpha=args.auto_scale_alpha, sdxl=args.sdxl)
-        ckpt_manager._save_ckpt(state, save_path=args.dump_path)
+        sd_unet = ckpt_loader.load(args.lora_path)
+        sd_TE = ckpt_loader.load(args.lora_path_TE) if args.lora_path_TE else {'base':{}}
+        state = converter.convert_to_webui(sd_unet['base'], sd_TE['base'], auto_scale_alpha=args.auto_scale_alpha, sdxl=args.sdxl)
+        ckpt_saver.save(state, args.dump_path)
         print('save lora to:', args.dump_path)

hcpdiff/tools/save_model.py

@@ -0,0 +1,12 @@
+from diffusers import DiffusionPipeline
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("model", default=None, type=str)
+parser.add_argument("output", default=None, type=str)
+args = parser.parse_args()
+
+pipe = DiffusionPipeline.from_pretrained(args.model, safety_checker=None, requires_safety_checker=False,
+                                        resume_download=True)
+
+pipe.save_pretrained(args.output)

hcpdiff/tools/sd2diffusers.py

@@ -211,7 +211,7 @@ def sd_vae_to_diffuser(args):
 def convert_ckpt(args):
     pipe = load_sd_ckpt(
         args.checkpoint_path,
-        original_config_file=args.original_config_file,
+        config_files={'v1': args.original_config_file},
         image_size=args.image_size,
         prediction_type=args.prediction_type,
         model_type=args.pipeline_type,

hcpdiff/train_colo.py

@@ -23,7 +23,7 @@ from colossalai.utils.model.colo_init_context import ColoInitContext
 from colossalai.utils.model.colo_init_context import _convert_to_coloparam
 from colossalai.tensor import ColoParameter
 
-from hcpdiff.train_ac import Trainer, get_scheduler, ModelEMA
+from hcpdiff.train_ac_old import Trainer, get_scheduler, ModelEMA
 from diffusers import UNet2DConditionModel
 from hcpdiff.utils.colo_utils import gemini_zero_dpp, GeminiAdamOptimizerP
 from hcpdiff.utils.utils import load_config_with_cli

hcpdiff/train_deepspeed.py

@@ -7,7 +7,7 @@ from functools import partial
 import torch
 
 from hcpdiff.ckpt_manager import CkptManagerPKL, CkptManagerSafe
-from hcpdiff.train_ac import Trainer, load_config_with_cli
+from hcpdiff.train_ac_old import Trainer, load_config_with_cli
 from hcpdiff.utils.net_utils import get_scheduler
 
 class TrainerDeepSpeed(Trainer):

hcpdiff/trainer_ac.py

@@ -0,0 +1,79 @@
+import argparse
+import warnings
+
+import torch
+from rainbowneko.parser import load_config_with_cli
+from rainbowneko.ckpt_manager import NekoSaver
+from rainbowneko.train import Trainer
+from rainbowneko.utils import xformers_available, is_dict
+from hcpdiff.ckpt_manager import EmbFormat
+
+class HCPTrainer(Trainer):
+    def config_model(self):
+        if self.cfgs.model.enable_xformers:
+            if xformers_available:
+                self.model_wrapper.enable_xformers()
+            else:
+                warnings.warn("xformers is not available. Make sure it is installed correctly")
+
+        if self.model_wrapper.vae is not None:
+            self.vae_dtype = self.weight_dtype_map.get(self.cfgs.model.get('vae_dtype', None), torch.float32)
+            self.model_wrapper.set_dtype(self.weight_dtype, self.vae_dtype)
+
+        if self.cfgs.model.gradient_checkpointing:
+            self.model_wrapper.enable_gradient_checkpointing()
+
+    def get_param_group_train(self):
+        train_params = super().get_param_group_train()
+
+        # For prompt-tuning
+        if self.cfgs.emb_pt is None:
+            train_params_emb, self.train_pts = [], {}
+        else:
+            from hcpdiff.parser import CfgEmbPTParser
+            self.cfgs.emb_pt: CfgEmbPTParser
+
+            train_params_emb, self.train_pts = self.cfgs.emb_pt.get_params_group(self.model_wrapper)
+            self.emb_format = EmbFormat()
+        train_params += train_params_emb
+        return train_params
+
+    @property
+    def pt_trainable(self):
+        return self.cfgs.emb_pt is not None
+
+    def get_loss(self, ds_name, model_pred, inputs):
+        loss = super().get_loss(ds_name, model_pred, inputs)
+        # make DDP happy
+        if len(self.train_pts)>0:
+            loss = loss+0*sum([emb.mean() for emb in self.train_pts.values()])
+        return loss
+
+    def save_model(self, from_raw=False):
+        NekoSaver.save_all(
+            self.model_raw,
+            plugin_groups={**self.all_plugin, 'embs': self.train_pts},
+            cfg=self.ckpt_saver,
+            model_ema=getattr(self, "ema_model", None),
+            name_template=f'{{}}-{self.real_step}',
+        )
+
+        self.loggers.info(f"Saved state, step: {self.real_step}")
+
+def hcp_train():
+    import subprocess
+    parser = argparse.ArgumentParser(description='HCP-Diffusion Launcher')
+    parser.add_argument('--launch_cfg', type=str, default='cfgs/launcher/multi.yaml')
+    args, train_args = parser.parse_known_args()
+
+    subprocess.run(["accelerate", "launch", '--config_file', args.launch_cfg, "-m",
+                    "hcpdiff.trainer_ac"] + train_args, check=True)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="HCP Diffusion Trainer")
+    parser.add_argument("--cfg", type=str, default=None, required=True)
+    args, cfg_args = parser.parse_known_args()
+
+    parser, conf = load_config_with_cli(args.cfg, args_list=cfg_args)  # skip --cfg
+    trainer = HCPTrainer(parser, conf)
+    trainer.train()

hcpdiff/trainer_ac_single.py

@@ -0,0 +1,31 @@
+import argparse
+import sys
+from functools import partial
+
+import torch
+from accelerate import Accelerator
+from loguru import logger
+
+from rainbowneko.train.trainer import TrainerSingleCard
+from hcpdiff.trainer_ac import HCPTrainer, load_config_with_cli
+
+class HCPTrainerSingleCard(TrainerSingleCard, HCPTrainer):
+    pass
+
+def hcp_train():
+    import subprocess
+    parser = argparse.ArgumentParser(description='HCP-Diffusion Launcher')
+    parser.add_argument('--launch_cfg', type=str, default='cfgs/launcher/single.yaml')
+    args, train_args = parser.parse_known_args()
+
+    subprocess.run(["accelerate", "launch", '--config_file', args.launch_cfg, "-m",
+                    "hcpdiff.trainer_ac_single"] + train_args, check=True)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='HCP Diffusion Trainer')
+    parser.add_argument("--cfg", type=str, default=None, required=True)
+    args, cfg_args = parser.parse_known_args()
+
+    parser, conf = load_config_with_cli(args.cfg, args_list=cfg_args)  # skip --cfg
+    trainer = HCPTrainerSingleCard(parser, conf)
+    trainer.train()

hcpdiff/utils/__init__.py

@@ -1,4 +1,2 @@
 from .utils import *
-from .img_size_tool import get_image_size
-from .cfg_resolvers import *
 from .net_utils import *

hcpdiff/utils/inpaint_pipe.py

@@ -21,18 +21,23 @@ import torch
 from packaging import version
 from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
 
+from diffusers import StableDiffusionInpaintPipelineLegacy
 from diffusers.configuration_utils import FrozenDict
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
 from diffusers.models import AutoencoderKL, UNet2DConditionModel
 from diffusers.models.lora import adjust_lora_scale_text_encoder
 from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import PIL_INTERPOLATION, USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils import PIL_INTERPOLATION, deprecate, logging
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 
+try:
+    from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers
+except:
+    USE_PEFT_BACKEND = False
 
 logger = logging.get_logger(__name__)
 

hcpdiff/utils/net_utils.py

@@ -6,11 +6,19 @@ import torch
 from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION, Optimizer
 from torch import nn
 from torch.optim import lr_scheduler
-from transformers import PretrainedConfig, AutoTokenizer
+from transformers import PretrainedConfig, AutoTokenizer, T5EncoderModel, CLIPTextModel
 from functools import partial
+from huggingface_hub import hf_hub_download
+import json
 
 dtype_dict = {'fp32':torch.float32, 'amp':torch.float32, 'fp16':torch.float16, 'bf16':torch.bfloat16}
 
+try:
+    dtype_dict['fp8_e4m3'] = torch.float8_e4m3fn
+    dtype_dict['fp8_e5m2'] = torch.float8_e5m2
+except:
+    pass
+
 def get_scheduler(cfg, optimizer):
     if cfg is None:
         return None
@@ -90,7 +98,7 @@ def auto_tokenizer_cls(pretrained_model_name_or_path: str, revision: str = None)
             revision=revision, use_fast=False,
         )
         return SDXLTokenizer
-    except OSError:
+    except:
         # not sdxl, only one tokenizer
         return AutoTokenizer
 
@@ -102,8 +110,10 @@ def auto_text_encoder_cls(pretrained_model_name_or_path: str, revision: str = No
             subfolder="text_encoder_2",
             revision=revision,
         )
+        if text_encoder_config.architectures is None:
+            raise ValueError()
         return SDXLTextEncoder
-    except OSError:
+    except:
         text_encoder_config = PretrainedConfig.from_pretrained(
             pretrained_model_name_or_path,
             subfolder="text_encoder",
@@ -112,16 +122,26 @@ def auto_text_encoder_cls(pretrained_model_name_or_path: str, revision: str = No
         model_class = text_encoder_config.architectures[0]
 
         if model_class == "CLIPTextModel":
-            from transformers import CLIPTextModel
-
             return CLIPTextModel
         elif model_class == "RobertaSeriesModelWithTransformation":
             from diffusers.pipelines.alt_diffusion.modeling_roberta_series import RobertaSeriesModelWithTransformation
 
             return RobertaSeriesModelWithTransformation
+        elif model_class == "T5EncoderModel":
+            return T5EncoderModel
         else:
             raise ValueError(f"{model_class} is not supported.")
 
+def get_pipe_name(path: str):
+    if os.path.isdir(path):
+        json_file = os.path.join(path, "model_index.json")
+    else:
+        json_file = hf_hub_download(path, "model_index.json")
+    with open(json_file, "r", encoding="utf-8") as reader:
+        text = reader.read()
+        data = json.loads(text)
+    return data['_class_name']
+
 def auto_tokenizer(pretrained_model_name_or_path: str, revision: str = None, **kwargs):
     return auto_tokenizer_cls(pretrained_model_name_or_path, revision).from_pretrained(pretrained_model_name_or_path, revision=revision, **kwargs)
 
@@ -225,4 +245,7 @@ def split_module_name(layer_name):
     return parent_name, host_name
 
 def get_dtype(dtype):
-    return dtype_dict.get(dtype, torch.float32)
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    else:
+        return dtype_dict.get(dtype, torch.float32)

hcpdiff/utils/pipe_hook.py

@@ -2,9 +2,9 @@ from typing import Union, List, Optional, Callable, Dict, Any
 
 import PIL
 import torch
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, PixArtTransformer2DModel
 from diffusers.image_processor import VaeImageProcessor
-from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
 from .inpaint_pipe import preprocess_mask, preprocess_image, StableDiffusionInpaintPipelineLegacy
 from einops import repeat
 
@@ -122,12 +122,20 @@ class HookPipe_T2I(StableDiffusionPipeline):
                 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
                 if pooled_output is None:
-                    noise_pred = self.unet(latent_model_input, t, prompt_embeds[i], encoder_attention_mask=encoder_attention_mask,
-                                           cross_attention_kwargs=cross_attention_kwargs, ).sample
+                    if isinstance(self.unet, PixArtTransformer2DModel):
+                        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+                        noise_pred = self.unet(latent_model_input, timestep=t.repeat(latent_model_input.shape[0]), encoder_hidden_states=prompt_embeds[i],
+                                            encoder_attention_mask=encoder_attention_mask,
+                                            cross_attention_kwargs=cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs).sample
+                    else:
+                        noise_pred = self.unet(latent_model_input, timestep=t, encoder_hidden_states=prompt_embeds[i],
+                                            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}
                     # predict the noise residual
-                    noise_pred = self.unet(latent_model_input, t, prompt_embeds[i], encoder_attention_mask=encoder_attention_mask,
+                    noise_pred = self.unet(latent_model_input, timestep=t, encoder_hidden_states=prompt_embeds[i],
+                                           encoder_attention_mask=encoder_attention_mask,
                                            cross_attention_kwargs=cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs).sample
 
                 # perform guidance
@@ -135,6 +143,10 @@ class HookPipe_T2I(StableDiffusionPipeline):
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                     noise_pred = noise_pred_uncond+guidance_scale*(noise_pred_text-noise_pred_uncond)
 
+                # learned sigma
+                if self.unet.config.out_channels // 2 == num_channels_latents:
+                    noise_pred = noise_pred.chunk(2, dim=1)[0]
+
                 # x_t -> x_0
                 alpha_prod_t = alphas_cumprod[t.long()]
                 beta_prod_t = 1-alpha_prod_t
@@ -271,8 +283,13 @@ class HookPipe_I2I(StableDiffusionImg2ImgPipeline):
 
                 # predict the noise residual
                 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 isinstance(self.unet, PixArtTransformer2DModel):
+                        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+                        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
+                    else:
+                        noise_pred = self.unet(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}
                     # predict the noise residual

hcpdiff/utils/utils.py

@@ -56,8 +56,8 @@ def remove_config_undefined(cfg):
 def load_config(path, remove_undefined=True):
     cfg = OmegaConf.load(path)
     if '_base_' in cfg:
-        for base in cfg['_base_']:
-            cfg = OmegaConf.merge(load_config(base, remove_undefined=False), cfg)
+        base_cfgs = [load_config(base, remove_undefined=False) for base in cfg['_base_']]
+        cfg = OmegaConf.merge(*base_cfgs, cfg)
         del cfg['_base_']
     if remove_undefined:
         cfg = remove_config_undefined(cfg)
@@ -85,7 +85,7 @@ def get_cfg_range(cfg_text:str):
 def to_validate_file(name):
     rstr = r"[\/\\\:\*\?\"\<\>\|]"  # '/ \ : * ? " < > |'
     new_title = re.sub(rstr, "_", name)  # 替换为下划线
-    return new_title[:50]
+    return new_title[:200]
 
 def make_mask(start, end, length):
     mask=torch.zeros(length)
@@ -159,4 +159,21 @@ def pad_attn_bias(x, attn_bias, block_size=8):
     # 在k维度上进行填充
     x_padded = F.pad(x, (0, 0, 0, padding_l, 0, 0), mode='constant', value=0)
     attn_bias_padded = F.pad(attn_bias, (0, padding_l, 0, 0), mode='constant', value=0)
-    return x_padded, attn_bias_padded
+    return x_padded, attn_bias_padded
+
+def linear_interp(t, x):
+    '''
+    t_l ---------t_h
+           ^x
+    '''
+    if (x>=len(t)).any():
+        x = x.clamp(max=len(t)-1e-6)
+    x0 = x.floor().long()
+    x1 = x0 + 1
+
+    y0 = t[x0]
+    y1 = t[x1]
+
+    xd = (x - x0.float())
+
+    return y0 * (1 - xd) + y1 * xd

hcpdiff/workflow/__init__.py

@@ -1,15 +1,20 @@
-from .base import BasicAction, MemoryMixin, from_memory, ExecAction, LoopAction
-from .diffusion import InputFeederAction, PrepareDiffusionAction, MakeLatentAction, NoisePredAction, SampleAction, DiffusionStepAction, \
-    X0PredAction, SeedAction, MakeTimestepsAction
+from .diffusion import InputFeederAction, MakeLatentAction, DenoiseAction, SampleAction, DiffusionStepAction, \
+    X0PredAction, SeedAction, MakeTimestepsAction, PrepareDiffusionAction, time_iter
 from .text import TextEncodeAction, TextHookAction, AttnMultTextEncodeAction
 from .vae import EncodeAction, DecodeAction
-from .io import LoadModelsAction, SaveImageAction, BuildModelLoaderAction, LoadPartAction, LoadLoraAction, LoadPluginAction
-from .utils import LatentResizeAction, ImageResizeAction
-from .model import VaeOptimizeAction, BuildOffloadAction, XformersEnableAction, StartTextEncode, StartDiffusion, EndTextEncode, EndDiffusion
+from .io import BuildModelsAction, SaveImageAction, LoadImageAction
+from .utils import LatentResizeAction, ImageResizeAction, FeedtoCNetAction
+from .model import VaeOptimizeAction, BuildOffloadAction, XformersEnableAction
+#from .flow import FilePromptAction
+
+try:
+    from .fast import SFastCompileAction
+except:
+    print('stable fast not installed.')
 
 from omegaconf import OmegaConf
 
-OmegaConf.register_new_resolver("hcp.from_memory", lambda mem_name: OmegaConf.create({
-    '_target_': 'hcpdiff.workflow.from_memory',
-    'mem_name': mem_name,
-}))
+OmegaConf.register_new_resolver("hcp.from_memory", lambda mem_name:OmegaConf.create({
+    '_target_':'hcpdiff.workflow.from_memory',
+    'mem_name':mem_name,
+}))