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

hcpdiff/models/textencoder_ex.py

@@ -8,29 +8,53 @@ textencoder_ex.py
     :Licence:     Apache-2.0
 """
 
-from typing import Tuple, Optional, List
+from typing import Tuple, Optional
 
 import torch
 from einops import repeat, rearrange
 from einops.layers.torch import Rearrange
+from loguru import logger
 from torch import nn
+from transformers import CLIPTextModelWithProjection, T5EncoderModel
 from transformers.models.clip.modeling_clip import CLIPAttention
 
 class TEEXHook:
-    def __init__(self, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, device='cuda', use_attention_mask=False):
+    def __init__(self, text_enc: nn.Module, tokenizer, N_repeats=1, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
         self.text_enc = text_enc
         self.tokenizer = tokenizer
 
         self.N_repeats = N_repeats
         self.clip_skip = clip_skip
         self.clip_final_norm = clip_final_norm
-        self.device = device
-        self.attn_mult = None
         self.use_attention_mask = use_attention_mask
 
         text_enc.register_forward_hook(self.forward_hook)
         text_enc.register_forward_pre_hook(self.forward_hook_input)
 
+    def find_final_norm(self, text_enc: nn.Module):
+        for module in text_enc.modules():
+            if 'final_layer_norm' in module._modules:
+                logger.info(f'find final_layer_norm in {type(module)}')
+                return module.final_layer_norm
+
+        logger.info(f'final_layer_norm not found in {type(text_enc)}')
+        return None
+
+    @property
+    def clip_final_norm(self):
+        return self.final_layer_norm is not None
+
+    @clip_final_norm.setter
+    def clip_final_norm(self, value: bool):
+        if value:
+            self.final_layer_norm = self.find_final_norm(self.text_enc)
+        else:
+            self.final_layer_norm = None
+
+    @property
+    def device(self):
+        return self.text_enc.device
+
     def encode_prompt_to_emb(self, prompt):
         text_inputs = self.tokenizer(
             prompt,
@@ -50,12 +74,23 @@ class TEEXHook:
         if position_ids is not None:
             position_ids = position_ids.to(self.device)
 
-        prompt_embeds, pooled_output = self.text_enc(
-            text_input_ids.to(self.device),
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            output_hidden_states=True,
-        )
+        # align with sd-webui
+        if isinstance(self.text_enc, CLIPTextModelWithProjection):
+            self.text_enc.text_projection.weight.data = self.text_enc.text_projection.weight.data.t()
+
+        if isinstance(self.text_enc, T5EncoderModel):
+            prompt_embeds, pooled_output = self.text_enc(
+                text_input_ids.to(self.device),
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        else:
+            prompt_embeds, pooled_output = self.text_enc(
+                text_input_ids.to(self.device),
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                output_hidden_states=True,
+            )
         return prompt_embeds, pooled_output, attention_mask
 
     def forward_hook_input(self, host, feat_in):
@@ -64,13 +99,12 @@ class TEEXHook:
 
     def forward_hook(self, host, feat_in: Tuple[torch.Tensor], feat_out):
         encoder_hidden_states = feat_out['hidden_states'][-self.clip_skip-1]
-        if self.clip_final_norm:
-            encoder_hidden_states = self.text_enc.text_model.final_layer_norm(encoder_hidden_states)
+        if self.clip_final_norm and self.final_layer_norm is not None:
+            encoder_hidden_states = self.final_layer_norm(encoder_hidden_states)
         if self.text_enc.training and self.clip_skip>0:
             encoder_hidden_states = encoder_hidden_states+0*feat_out['last_hidden_state'].mean()  # avoid unused parameters, make gradient checkpointing happy
-
         encoder_hidden_states = rearrange(encoder_hidden_states, '(b r) ... -> b r ...', r=self.N_repeats)  # [B, N_repeat, N_word+2, N_emb]
-        pooled_output = feat_out.pooler_output
+        pooled_output = feat_out.get('pooler_output', feat_out.get('text_embeds', None))
         # TODO: may have better fusion method
         if pooled_output is not None:
             pooled_output = rearrange(pooled_output, '(b r) ... -> b r ...', r=self.N_repeats).mean(dim=1)
@@ -81,7 +115,7 @@ class TEEXHook:
         return encoder_hidden_states, pooled_output
 
     def pool_hidden_states(self, encoder_hidden_states, input_ids):
-        pooled_output = encoder_hidden_states[:, :, -1, :].mean(dim=1) # [B, N_emb]
+        pooled_output = encoder_hidden_states[:, :, -1, :].mean(dim=1)  # [B, N_emb]
         return pooled_output
 
     @staticmethod
@@ -147,9 +181,11 @@ class TEEXHook:
         layer.forward = forward
 
     @classmethod
-    def hook(cls, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, device='cuda', use_attention_mask=False):
-        return cls(text_enc, tokenizer, N_repeats=N_repeats, clip_skip=clip_skip, clip_final_norm=clip_final_norm, device=device, use_attention_mask=use_attention_mask)
+    def hook(cls, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
+        return cls(text_enc, tokenizer, N_repeats=N_repeats, clip_skip=clip_skip, clip_final_norm=clip_final_norm,
+                   use_attention_mask=use_attention_mask)
 
     @classmethod
     def hook_pipe(cls, pipe, N_repeats=3, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
-        return cls(pipe.text_encoder, pipe.tokenizer, N_repeats=N_repeats, device='cuda', clip_skip=clip_skip, clip_final_norm=clip_final_norm, use_attention_mask=use_attention_mask)
+        return cls(pipe.text_encoder, pipe.tokenizer, N_repeats=N_repeats, clip_skip=clip_skip, clip_final_norm=clip_final_norm,
+                   use_attention_mask=use_attention_mask)

hcpdiff/models/wrapper/__init__.py

@@ -0,0 +1,3 @@
+from .sd import SD15Wrapper, SDXLWrapper
+from .pixart import PixArtWrapper
+from .utils import TEHookCFG, SD15_TEHookCFG, SDXL_TEHookCFG

hcpdiff/models/wrapper/pixart.py

@@ -0,0 +1,19 @@
+from .sd import SD15Wrapper
+from hcpdiff.utils import pad_attn_bias
+
+class PixArtWrapper(SD15Wrapper):
+    def forward_denoiser(self, x_t, prompt_ids, encoder_hidden_states, timesteps, attn_mask=None, position_ids=None, resolution=None, aspect_ratio=None,
+                     plugin_input={}, **kwargs):
+        if attn_mask is not None:
+            attn_mask[:, :self.min_attnmask] = 1
+            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask,
+                         encoder_hidden_states=encoder_hidden_states, **plugin_input)
+        if hasattr(self.denoiser, 'input_feeder'):
+            for feeder in self.denoiser.input_feeder:
+                feeder(input_all)
+        added_cond_kwargs = {"resolution":resolution, "aspect_ratio":aspect_ratio}
+        model_pred = self.denoiser(x_t, encoder_hidden_states, timesteps, encoder_attention_mask=attn_mask,
+                               added_cond_kwargs=added_cond_kwargs).sample  # Predict the noise residual
+        return model_pred

hcpdiff/models/wrapper/sd.py

@@ -0,0 +1,218 @@
+from contextlib import nullcontext
+from functools import partial
+from typing import Dict, Union
+
+import torch
+from diffusers import AutoencoderKL, UNet2DConditionModel
+from rainbowneko.models.wrapper import BaseWrapper
+from torch import Tensor
+from torch import nn
+
+from hcpdiff.diffusion.sampler import BaseSampler
+from hcpdiff.models import TEEXHook
+from hcpdiff.models.compose import ComposeTEEXHook
+from hcpdiff.utils import pad_attn_bias
+from .utils import TEHookCFG, SD15_TEHookCFG, SDXL_TEHookCFG
+from ..cfg_context import CFGContext
+
+class SD15Wrapper(BaseWrapper):
+    def __init__(self, denoiser: UNet2DConditionModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
+                 pred_type='eps', TE_hook_cfg:TEHookCFG=SD15_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
+        super().__init__()
+        self.key_mapper_in = self.build_mapper(key_map_in, None, (
+            'prompt -> prompt_ids', 'image -> image', 'attn_mask -> attn_mask', 'position_ids -> position_ids', 'neg_prompt -> neg_prompt_ids',
+            'neg_attn_mask -> neg_attn_mask', 'neg_position_ids -> neg_position_ids', 'plugin_input -> plugin_input'))
+        self.key_mapper_out = self.build_mapper(key_map_out, None, None)
+
+        self.denoiser = denoiser
+        self.TE = TE
+        self.vae = vae
+        self.noise_sampler = noise_sampler
+        self.tokenizer = tokenizer
+        self.min_attnmask = min_attnmask
+
+        self.pred_type = pred_type
+
+        self.TE_hook_cfg = TEHookCFG.create(TE_hook_cfg)
+        self.cfg_context = cfg_context
+        self.tokenizer.N_repeats = self.TE_hook_cfg.tokenizer_repeats
+
+    def post_init(self):
+        self.make_TE_hook(self.TE_hook_cfg)
+
+        self.vae_trainable = False
+        if self.vae is not None:
+            for p in self.vae.parameters():
+                if p.requires_grad:
+                    self.vae_trainable = True
+                    break
+
+        self.TE_trainable = False
+        for p in self.TE.parameters():
+            if p.requires_grad:
+                self.TE_trainable = True
+                break
+
+    def make_TE_hook(self, TE_hook_cfg):
+        # Hook and extend text_encoder
+        self.text_enc_hook = TEEXHook.hook(self.TE, self.tokenizer, N_repeats=TE_hook_cfg.tokenizer_repeats,
+                                           clip_skip=TE_hook_cfg.clip_skip, clip_final_norm=TE_hook_cfg.clip_final_norm)
+
+    def get_latents(self, image: Tensor):
+        if image.shape[1] == 3:
+            with torch.no_grad() if self.vae_trainable else nullcontext():
+                latents = self.vae.encode(image.to(dtype=self.vae.dtype)).latent_dist.sample()
+                latents = latents*self.vae.config.scaling_factor
+        else:
+            latents = image  # Cached latents
+        return latents
+
+    def forward_TE(self, prompt_ids, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask, **plugin_input)
+        if hasattr(self.TE, 'input_feeder'):
+            for feeder in self.TE.input_feeder:
+                feeder(input_all)
+        # Get the text embedding for conditioning
+        encoder_hidden_states = self.TE(prompt_ids, position_ids=position_ids, attention_mask=attn_mask, output_hidden_states=True)[0]
+        return encoder_hidden_states
+
+    def forward_denoiser(self, x_t, prompt_ids, encoder_hidden_states, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
+        if attn_mask is not None:
+            attn_mask[:, :self.min_attnmask] = 1
+            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask,
+                         encoder_hidden_states=encoder_hidden_states, **plugin_input)
+        if hasattr(self.denoiser, 'input_feeder'):
+            for feeder in self.denoiser.input_feeder:
+                feeder(input_all)
+        model_pred = self.denoiser(x_t, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask).sample  # Predict the noise residual
+        return model_pred
+
+    def model_forward(self, prompt_ids, image, attn_mask=None, position_ids=None, neg_prompt_ids=None, neg_attn_mask=None, neg_position_ids=None,
+                      plugin_input={}, **kwargs):
+        # input prepare
+        x_0 = self.get_latents(image)
+        x_t, noise, sigma, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
+        x_t_in = x_t*self.noise_sampler.c_in(sigma).to(dtype=x_t.dtype)
+
+        if neg_prompt_ids:
+            prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
+            if neg_attn_mask:
+                attn_mask = torch.cat([neg_attn_mask, attn_mask], dim=0)
+            if neg_position_ids:
+                position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
+
+        # model forward
+        x_t_in, timesteps = self.cfg_context.pre(x_t_in, timesteps)
+        encoder_hidden_states = self.forward_TE(prompt_ids, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+                                                plugin_input=plugin_input, **kwargs)
+        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+                                           plugin_input=plugin_input, **kwargs)
+        model_pred = self.cfg_context.post(model_pred)
+
+        return dict(model_pred=model_pred, noise=noise, sigma=sigma, timesteps=timesteps, x_0=x_0, x_t=x_t, pred_type=self.pred_type,
+                    noise_sampler=self.noise_sampler)
+
+    def forward(self, ds_name=None, **kwargs):
+        model_args, model_kwargs = self.get_map_data(self.key_mapper_in, kwargs, ds_name)
+        out = self.model_forward(*model_args, **model_kwargs)
+        return self.get_map_data(self.key_mapper_out, out, ds_name=ds_name)[1]
+
+    def enable_gradient_checkpointing(self):
+        def grad_ckpt_enable(m):
+            if getattr(m, 'gradient_checkpointing', False):
+                m.training = True
+
+        self.denoiser.enable_gradient_checkpointing()
+        if self.TE_trainable:
+            self.TE.gradient_checkpointing_enable()
+        self.apply(grad_ckpt_enable)
+
+    def enable_xformers(self):
+        self.denoiser.enable_xformers_memory_efficient_attention()
+
+    @property
+    def trainable_parameters(self):
+        return [p for p in self.parameters() if p.requires_grad]
+
+    @property
+    def trainable_models(self) -> Dict[str, nn.Module]:
+        return {'self':self}
+
+    def set_dtype(self, dtype, vae_dtype):
+        self.dtype = dtype
+        self.vae_dtype = vae_dtype
+        # Move vae and text_encoder to device and cast to weight_dtype
+        if self.vae is not None:
+            self.vae = self.vae.to(dtype=vae_dtype)
+        if not self.TE_trainable:
+            self.TE = self.TE.to(dtype=dtype)
+
+    @classmethod
+    def from_pretrained(cls, models: Union[partial, Dict[str, nn.Module]], **kwargs):
+        models = models() if isinstance(models, partial) else models
+        return cls(models['denoiser'], models['TE'], models['vae'], models['noise_sampler'], models['tokenizer'], **kwargs)
+
+class SDXLWrapper(SD15Wrapper):
+    def __init__(self, denoiser: UNet2DConditionModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
+                 pred_type='eps', TE_hook_cfg:TEHookCFG=SDXL_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
+        super().__init__(denoiser, TE, vae, noise_sampler, tokenizer, min_attnmask, pred_type, TE_hook_cfg, cfg_context, key_map_in, key_map_out)
+        self.key_mapper_in = self.build_mapper(key_map_in, None, (
+            'prompt -> prompt_ids', 'image -> image', 'attn_mask -> attn_mask', 'position_ids -> position_ids', 'neg_prompt -> neg_prompt_ids',
+            'neg_attn_mask -> neg_attn_mask', 'neg_position_ids -> neg_position_ids', 'plugin_input -> plugin_input', 'coord -> crop_info'))
+
+    def make_TE_hook(self, TE_hook_cfg):
+        # Hook and extend text_encoder
+        self.text_enc_hook = ComposeTEEXHook.hook(self.TE, self.tokenizer, N_repeats=TE_hook_cfg.tokenizer_repeats,
+                                                  clip_skip=TE_hook_cfg.clip_skip, clip_final_norm=TE_hook_cfg.clip_final_norm)
+
+    def forward_TE(self, prompt_ids, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask, **plugin_input)
+        if hasattr(self.TE, 'input_feeder'):
+            for feeder in self.TE.input_feeder:
+                feeder(input_all)
+        # Get the text embedding for conditioning
+        encoder_hidden_states, pooled_output = self.TE(prompt_ids, position_ids=position_ids, attention_mask=attn_mask, output_hidden_states=True)
+        return encoder_hidden_states, pooled_output
+
+    def forward_denoiser(self, x_t, prompt_ids, encoder_hidden_states, timesteps, added_cond_kwargs, attn_mask=None, position_ids=None,
+                         plugin_input={}, **kwargs):
+        if attn_mask is not None:
+            attn_mask[:, :self.min_attnmask] = 1
+            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask,
+                         encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs, **plugin_input)
+        if hasattr(self.denoiser, 'input_feeder'):
+            for feeder in self.denoiser.input_feeder:
+                feeder(input_all)
+        model_pred = self.denoiser(x_t, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask,
+                                   added_cond_kwargs=added_cond_kwargs).sample  # Predict the noise residual
+        return model_pred
+
+    def model_forward(self, prompt_ids, image, attn_mask=None, position_ids=None, neg_prompt_ids=None, neg_attn_mask=None, neg_position_ids=None,
+                      crop_info=None, plugin_input={}):
+        # input prepare
+        x_0 = self.get_latents(image)
+        x_t, noise, sigma, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
+        x_t_in = x_t*self.noise_sampler.c_in(sigma).to(dtype=x_t.dtype)
+
+        if neg_prompt_ids:
+            prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
+            if neg_attn_mask:
+                attn_mask = torch.cat([neg_attn_mask, attn_mask], dim=0)
+            if neg_position_ids:
+                position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
+
+        # model forward
+        x_t_in, timesteps = self.cfg_context.pre(x_t_in, timesteps)
+        encoder_hidden_states, pooled_output = self.forward_TE(prompt_ids, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+                                                               plugin_input=plugin_input)
+        added_cond_kwargs = {"text_embeds":pooled_output[-1], "time_ids":crop_info}
+        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, timesteps, added_cond_kwargs=added_cond_kwargs,
+                                           attn_mask=attn_mask, position_ids=position_ids, plugin_input=plugin_input)
+        model_pred = self.cfg_context.post(model_pred)
+
+        return dict(model_pred=model_pred, noise=noise, sigma=sigma, timesteps=timesteps, x_0=x_0, x_t=x_t, pred_type=self.pred_type,
+                    noise_sampler=self.noise_sampler)

hcpdiff/models/wrapper/utils.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from rainbowneko.utils import is_dict
+
+class TEHookCFG:
+    def __init__(self, tokenizer_repeats: int = 1, clip_skip: int = 0, clip_final_norm: bool = True):
+        self.tokenizer_repeats = tokenizer_repeats
+        self.clip_skip = clip_skip
+        self.clip_final_norm = clip_final_norm
+
+    @classmethod
+    def create(cls, cfg):
+        if is_dict(cfg):
+            return cls(**cfg)
+        elif isinstance(cfg, cls):
+            return cfg
+        else:
+            raise ValueError(f'Invalid TEHookCFG type: {type(cfg)}')
+        
+SD15_TEHookCFG = TEHookCFG()
+SDXL_TEHookCFG = TEHookCFG(clip_skip=1, clip_final_norm=False)

hcpdiff/parser/__init__.py

@@ -0,0 +1 @@
+from .embpt import CfgEmbPTParser

hcpdiff/parser/embpt.py

@@ -0,0 +1,32 @@
+from typing import Dict, Tuple, List
+from rainbowneko.utils import Path_Like
+from hcpdiff.models import EmbeddingPTHook
+from torch import Tensor
+
+class CfgEmbPTParser:
+    def __init__(self, emb_dir: Path_Like, cfg_pt: Dict[str, Dict], lr: float = 1e-5, weight_decay: float = 0):
+        self.emb_dir = emb_dir
+        self.cfg_pt = cfg_pt
+        self.lr = lr
+        self.weight_decay = weight_decay
+
+    def get_params_group(self, model) -> Tuple[List, Dict[str, Tensor]]:
+        self.embedding_hook, self.ex_words_emb = EmbeddingPTHook.hook_from_dir(
+            self.emb_dir, model.tokenizer, model.TE, N_repeats=model.tokenizer.N_repeats)
+        self.embedding_hook.requires_grad_(False)
+
+        train_params_emb = []
+        train_pts = {}
+        for pt_name, info in self.cfg_pt.items():
+            word_emb = self.ex_words_emb[pt_name]
+            train_pts[pt_name] = word_emb
+            word_emb.requires_grad = True
+            self.embedding_hook.emb_train.append(word_emb)
+            param_group = {'params':word_emb}
+            if 'lr' in info:
+                param_group['lr'] = info.lr
+            if 'weight_decay' in info:
+                param_group['weight_decay'] = info.weight_decay
+            train_params_emb.append(param_group)
+
+        return train_params_emb, train_pts

hcpdiff/tools/convert_caption_txt2json.py

@@ -2,7 +2,7 @@ import argparse
 import json
 import os
 
-from hcpdiff.utils.img_size_tool import types_support
+from rainbowneko.utils import types_support
 
 parser = argparse.ArgumentParser(description='Stable Diffusion Training')
 parser.add_argument('--data_root', type=str, default='')

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')