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

hcpdiff/utils/img_size_tool.py

@@ -1,248 +0,0 @@
-"""
-
-img_size_tool.py
-====================
-
-    :Name:        get_image_size
-    :Purpose:     extract image dimensions given a file path
-
-    :Author:      Paulo Scardine (based on code from Emmanuel VAÏSSE)
-                  Dong Ziyi, add webp support
-
-    :Created:     26/09/2013
-    :Modified:    02/03/2023
-    :Copyright:   (c) Paulo Scardine 2013
-    :Licence:     MIT
-
-"""
-
-import collections
-import os
-import io
-import struct
-from PIL import Image
-
-FILE_UNKNOWN = "Sorry, don't know how to get size for this file."
-
-class UnknownImageFormat(Exception):
-    pass
-
-types_support = ['bmp', 'gif', 'ico', 'jpeg', 'jpg', 'png', 'tiff', 'webp']
-
-def get_image_size(file_path):
-    """
-    Return (width, height) for a given img file content - no external
-    dependencies except the os and struct builtin modules
-    """
-    width, height = get_image_metadata(file_path)
-    return width, height
-
-
-def get_image_size_from_bytesio(input, size):
-    """
-    Return (width, height) for a given img file content - no external
-    dependencies except the os and struct builtin modules
-
-    Args:
-        input (io.IOBase): io object support read & seek
-        size (int): size of buffer in byte
-    """
-    width, height = get_image_metadata_from_bytesio(input, size)
-    return width, height
-
-
-def get_image_metadata(file_path):
-    """
-    Return an `Image` object for a given img file content - no external
-    dependencies except the os and struct builtin modules
-
-    Args:
-        file_path (str): path to an image file
-
-    Returns:
-        (width, height)
-    """
-    size = os.path.getsize(file_path)
-
-    # be explicit with open arguments - we need binary mode
-    with io.open(file_path, "rb") as input:
-        return get_image_metadata_from_bytesio(input, size, file_path)
-
-
-def get_image_metadata_from_bytesio(input, size, file_path=None):
-    """
-    Return an `Image` object for a given img file content - no external
-    dependencies except the os and struct builtin modules
-
-    Args:
-        input (io.IOBase): io object support read & seek
-        size (int): size of buffer in byte
-        file_path (str): path to an image file
-
-    Returns:
-        (width, height)
-    """
-    height = -1
-    width = -1
-    data = input.read(30)
-    msg = " raised while trying to decode as JPEG."
-
-    if (size >= 10) and data[:6] in (b'GIF87a', b'GIF89a'):
-        # GIFs
-        #imgtype = GIF
-        w, h = struct.unpack("<HH", data[6:10])
-        width = int(w)
-        height = int(h)
-    elif (size >= 24) and data[8:12] == b'WEBP':
-        # WEBPs
-        #imgtype = WEBP
-        if data[15]==b'X': #VP8X
-            w = int.from_bytes(data[24:27], 'little')+1
-            h = int.from_bytes(data[27:30], 'little')+1
-        elif data[15]==b' ': #VP8
-            w, h = struct.unpack("<HH", data[0x1A:0x1E])
-        else:
-            w, h = Image.open(file_path).size
-
-        width = int(w)
-        height = int(h)
-    elif ((size >= 24) and data.startswith(b'\211PNG\r\n\032\n')
-            and (data[12:16] == b'IHDR')):
-        # PNGs
-        #imgtype = PNG
-        w, h = struct.unpack(">LL", data[16:24])
-        width = int(w)
-        height = int(h)
-    elif (size >= 16) and data.startswith(b'\211PNG\r\n\032\n'):
-        # older PNGs
-        #imgtype = PNG
-        w, h = struct.unpack(">LL", data[8:16])
-        width = int(w)
-        height = int(h)
-    elif (size >= 2) and data.startswith(b'\377\330'):
-        # JPEG
-        #imgtype = JPEG
-        input.seek(0)
-        input.read(2)
-        b = input.read(1)
-        try:
-            while (b and ord(b) != 0xDA):
-                while (ord(b) != 0xFF):
-                    b = input.read(1)
-                while (ord(b) == 0xFF):
-                    b = input.read(1)
-                if (ord(b) >= 0xC0 and ord(b) <= 0xC3):
-                    input.read(3)
-                    h, w = struct.unpack(">HH", input.read(4))
-                    break
-                else:
-                    input.read(
-                        int(struct.unpack(">H", input.read(2))[0]) - 2)
-                b = input.read(1)
-            width = int(w)
-            height = int(h)
-        except struct.error:
-            raise UnknownImageFormat("StructError" + msg)
-        except ValueError:
-            raise UnknownImageFormat("ValueError" + msg)
-        except Exception as e:
-            raise UnknownImageFormat(e.__class__.__name__ + msg)
-    elif (size >= 26) and data.startswith(b'BM'):
-        # BMP
-        #imgtype = BMP
-        headersize = struct.unpack("<I", data[14:18])[0]
-        if headersize == 12:
-            w, h = struct.unpack("<HH", data[18:22])
-            width = int(w)
-            height = int(h)
-        elif headersize >= 40:
-            w, h = struct.unpack("<ii", data[18:26])
-            width = int(w)
-            # as h is negative when stored upside down
-            height = abs(int(h))
-        else:
-            raise UnknownImageFormat(
-                "Unkown DIB header size:" +
-                str(headersize))
-    elif (size >= 8) and data[:4] in (b"II\052\000", b"MM\000\052"):
-        # Standard TIFF, big- or little-endian
-        # BigTIFF and other different but TIFF-like formats are not
-        # supported currently
-        #imgtype = TIFF
-        byteOrder = data[:2]
-        boChar = ">" if byteOrder == "MM" else "<"
-        # maps TIFF type id to size (in bytes)
-        # and python format char for struct
-        tiffTypes = {
-            1: (1, boChar + "B"),  # BYTE
-            2: (1, boChar + "c"),  # ASCII
-            3: (2, boChar + "H"),  # SHORT
-            4: (4, boChar + "L"),  # LONG
-            5: (8, boChar + "LL"),  # RATIONAL
-            6: (1, boChar + "b"),  # SBYTE
-            7: (1, boChar + "c"),  # UNDEFINED
-            8: (2, boChar + "h"),  # SSHORT
-            9: (4, boChar + "l"),  # SLONG
-            10: (8, boChar + "ll"),  # SRATIONAL
-            11: (4, boChar + "f"),  # FLOAT
-            12: (8, boChar + "d")   # DOUBLE
-        }
-        ifdOffset = struct.unpack(boChar + "L", data[4:8])[0]
-        try:
-            countSize = 2
-            input.seek(ifdOffset)
-            ec = input.read(countSize)
-            ifdEntryCount = struct.unpack(boChar + "H", ec)[0]
-            # 2 bytes: TagId + 2 bytes: type + 4 bytes: count of values + 4
-            # bytes: value offset
-            ifdEntrySize = 12
-            for i in range(ifdEntryCount):
-                entryOffset = ifdOffset + countSize + i * ifdEntrySize
-                input.seek(entryOffset)
-                tag = input.read(2)
-                tag = struct.unpack(boChar + "H", tag)[0]
-                if(tag == 256 or tag == 257):
-                    # if type indicates that value fits into 4 bytes, value
-                    # offset is not an offset but value itself
-                    type = input.read(2)
-                    type = struct.unpack(boChar + "H", type)[0]
-                    if type not in tiffTypes:
-                        raise UnknownImageFormat(
-                            "Unkown TIFF field type:" +
-                            str(type))
-                    typeSize = tiffTypes[type][0]
-                    typeChar = tiffTypes[type][1]
-                    input.seek(entryOffset + 8)
-                    value = input.read(typeSize)
-                    value = int(struct.unpack(typeChar, value)[0])
-                    if tag == 256:
-                        width = value
-                    else:
-                        height = value
-                if width > -1 and height > -1:
-                    break
-        except Exception as e:
-            raise UnknownImageFormat(str(e))
-    elif size >= 2:
-            # see http://en.wikipedia.org/wiki/ICO_(file_format)
-        #imgtype = 'ICO'
-        input.seek(0)
-        reserved = input.read(2)
-        if 0 != struct.unpack("<H", reserved)[0]:
-            raise UnknownImageFormat(FILE_UNKNOWN)
-        format = input.read(2)
-        assert 1 == struct.unpack("<H", format)[0]
-        num = input.read(2)
-        num = struct.unpack("<H", num)[0]
-        if num > 1:
-            import warnings
-            warnings.warn("ICO File contains more than one image")
-        # http://msdn.microsoft.com/en-us/library/ms997538.aspx
-        w = input.read(1)
-        h = input.read(1)
-        width = ord(w)
-        height = ord(h)
-    else:
-        raise UnknownImageFormat(FILE_UNKNOWN)
-
-    return width, height

hcpdiff/vis/__init__.py

@@ -1,3 +0,0 @@
-from .base_interface import BaseInterface
-from .disk_interface import DiskInterface
-from .webui_interface import WebUIInterface

hcpdiff/vis/base_interface.py

@@ -1,12 +0,0 @@
-
-class BaseInterface:
-    need_inter_imgs = False
-
-    def __init__(self, show_steps=0):
-        self.show_steps = show_steps
-
-    def on_inter_step(self, i, num_steps, t, latents, images):
-        pass
-
-    def on_infer_finish(self, images, prompt, negative_prompt, save_cfg=False, seeds=None):
-        pass

hcpdiff/vis/disk_interface.py

@@ -1,48 +0,0 @@
-import os
-
-from hcpdiff.utils.img_size_tool import types_support
-from hcpdiff.utils.utils import to_validate_file
-from omegaconf import OmegaConf
-
-from .base_interface import BaseInterface
-
-class DiskInterface(BaseInterface):
-    def __init__(self, save_root, save_cfg=True, image_type='png', quality=95, show_steps=0):
-        super(DiskInterface, self).__init__(show_steps=show_steps)
-        os.makedirs(save_root, exist_ok=True)
-        self.save_root = save_root
-        self.save_cfg = save_cfg
-        self.image_type = image_type
-        self.quality = quality
-
-        self.inter_imgs = []
-        if show_steps>0:
-            self.need_inter_imgs = True
-
-    def on_inter_step(self, i, num_steps, t, latents, images):
-        if len(self.inter_imgs) == 0:
-            for _ in range(len(images)):
-                self.inter_imgs.append([])
-        for u, img in enumerate(images):
-            self.inter_imgs[u].append(img)
-
-    def on_save_one(self, num_img_exist, img_path):
-        pass
-
-    def on_infer_finish(self, images, prompt, negative_prompt, cfgs_raw=None, seeds=None):
-        num_img_exist = max([0]+[int(x.split('-', 1)[0]) for x in os.listdir(self.save_root) if x.rsplit('.', 1)[-1] in types_support])+1
-
-        for bid, (p, pn, img) in enumerate(zip(prompt, negative_prompt, images)):
-            img_path = os.path.join(self.save_root, f"{num_img_exist}-{seeds[bid]}-{to_validate_file(prompt[0])}.{self.image_type}")
-            img.save(img_path, quality=self.quality)
-            self.on_save_one(num_img_exist, img_path)
-
-            if self.save_cfg and cfgs_raw is not None:
-                with open(os.path.join(self.save_root, f"{num_img_exist}-{seeds[bid]}-info.yaml"), 'w', encoding='utf-8') as f:
-                    cfgs_raw.seed = seeds[bid]
-                    f.write(OmegaConf.to_yaml(cfgs_raw))
-            if self.need_inter_imgs:
-                inter = self.inter_imgs[bid]
-                inter[0].save(os.path.join(self.save_root, f'{num_img_exist}-{seeds[bid]}-steps.webp'), "webp", save_all=True,
-                              append_images=inter[1:], duration=100)
-            num_img_exist += 1

hcpdiff/vis/webui_interface.py

@@ -1,17 +0,0 @@
-from .disk_interface import DiskInterface
-from loguru import logger
-
-class WebUIInterface(DiskInterface):
-
-    def __init__(self, save_root, image_type='png', quality=95, show_steps=1, show_inter=False):
-        super(WebUIInterface, self).__init__(save_root, image_type, quality, show_steps)
-        self.show_inter = show_inter
-        self.need_inter_imgs = self.need_inter_imgs and show_inter
-
-    def on_inter_step(self, i, num_steps, t, latents, images):
-        if self.show_inter:
-            super(WebUIInterface, self).on_inter_step(i, num_steps, t, latents, images)
-        logger.info(f'\nthis progress steps: {i}/{num_steps}')
-
-    def on_save_one(self, num_img_exist, img_path):
-        logger.info(f'this images output path: {img_path}')

hcpdiff/viser_fast.py

@@ -1,138 +0,0 @@
-import argparse
-import os
-from typing import List
-
-import hydra
-import torch
-from sfast.compilers.diffusion_pipeline_compiler import (compile_unet, CompilationConfig)
-from torch.cuda.amp import autocast
-
-from hcpdiff import Visualizer
-from hcpdiff.models import TokenizerHook
-from hcpdiff.models.compose import ComposeTEEXHook, ComposeEmbPTHook, ComposeTextEncoder
-from hcpdiff.utils.net_utils import to_cuda
-from hcpdiff.utils.utils import load_config_with_cli, prepare_seed, is_list, pad_attn_bias
-
-class VisualizerFast(Visualizer):
-    dtype_dict = {'fp32':torch.float32, 'fp16':torch.float16, 'bf16':torch.bfloat16}
-
-    def __init__(self, cfgs):
-        self.cfgs_raw = cfgs
-        self.cfgs = hydra.utils.instantiate(self.cfgs_raw)
-        self.cfg_merge = self.cfgs.merge
-        self.offload = 'offload' in self.cfgs and self.cfgs.offload is not None
-        self.dtype = self.dtype_dict[self.cfgs.dtype]
-
-        self.need_inter_imgs = any(item.need_inter_imgs for item in self.cfgs.interface)
-
-        self.pipe = self.load_model(self.cfgs.pretrained_model)
-
-        if self.cfg_merge:
-            self.merge_model()
-
-        # self.pipe = self.pipe.to(torch_dtype=self.dtype)
-
-        if isinstance(self.pipe.text_encoder, ComposeTextEncoder):
-            self.pipe.vae = self.pipe.vae.to(dtype=torch.float32)
-
-        if 'save_model' in self.cfgs and self.cfgs.save_model is not None:
-            self.save_model(self.cfgs.save_model)
-            os._exit(0)
-
-        self.build_optimize()
-        self.compile_model()
-
-    def build_optimize(self):
-        if self.offload:
-            self.build_offload(self.cfgs.offload)
-        else:
-            self.pipe.unet.to('cuda')
-        self.build_vae_offload()
-
-        if getattr(self.cfgs, 'vae_optimize', None) is not None:
-            if self.cfgs.vae_optimize.tiling:
-                self.pipe.vae.enable_tiling()
-            if self.cfgs.vae_optimize.slicing:
-                self.pipe.vae.enable_slicing()
-
-        self.emb_hook, _ = ComposeEmbPTHook.hook_from_dir(self.cfgs.emb_dir, self.pipe.tokenizer, self.pipe.text_encoder,
-                                                          N_repeats=self.cfgs.N_repeats)
-        self.te_hook = ComposeTEEXHook.hook_pipe(self.pipe, N_repeats=self.cfgs.N_repeats, clip_skip=self.cfgs.clip_skip,
-                                                 clip_final_norm=self.cfgs.clip_final_norm, use_attention_mask=self.cfgs.encoder_attention_mask)
-        self.token_ex = TokenizerHook(self.pipe.tokenizer)
-
-    def compile_model(self):
-        # compile model
-        config = CompilationConfig.Default()
-        config.enable_xformers = False
-        try:
-            import xformers
-            config.enable_xformers = True
-        except ImportError:
-            print('xformers not installed, skip')
-        # NOTE:
-        # When GPU VRAM is insufficient or the architecture is too old, Triton might be slow.
-        # Disable Triton if you encounter this problem.
-        try:
-            import tritonx
-            config.enable_triton = True
-        except ImportError:
-            print('Triton not installed, skip')
-        config.enable_cuda_graph = True
-
-        self.pipe.unet = compile_unet(self.pipe.unet, config)
-
-    @torch.inference_mode()
-    def vis_images(self, prompt, negative_prompt='', seeds: List[int] = None, **kwargs):
-        G = prepare_seed(seeds or [None]*len(prompt))
-
-        ex_input_dict, pipe_input_dict = self.get_ex_input()
-        kwargs.update(pipe_input_dict)
-
-        to_cuda(self.pipe.text_encoder)
-
-        mult_p, clean_text_p = self.token_ex.parse_attn_mult(prompt)
-        mult_n, clean_text_n = self.token_ex.parse_attn_mult(negative_prompt)
-        with autocast(enabled=self.cfgs.amp, dtype=self.dtype):
-            emb, pooled_output, attention_mask = self.te_hook.encode_prompt_to_emb(clean_text_n+clean_text_p)
-            if self.cfgs.encoder_attention_mask:
-                emb, attention_mask = pad_attn_bias(emb, attention_mask)
-            else:
-                attention_mask = None
-            emb_n, emb_p = emb.chunk(2)
-            emb_p = self.te_hook.mult_attn(emb_p, mult_p)
-            emb_n = self.te_hook.mult_attn(emb_n, mult_n)
-
-            # to_cpu(self.pipe.text_encoder)
-            # to_cuda(self.pipe.unet)
-
-            if hasattr(self.pipe.unet, 'input_feeder'):
-                for feeder in self.pipe.unet.input_feeder:
-                    feeder(ex_input_dict)
-
-            images = self.pipe(prompt_embeds=emb_p, negative_prompt_embeds=emb_n, callback=None, generator=G,
-                               pooled_output=pooled_output[-1], encoder_attention_mask=attention_mask, **kwargs).images
-        return images
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Fast HCP Diffusion Inference')
-    parser.add_argument('--cfg', type=str, default='cfgs/infer/text2img.yaml')
-    args, cfg_args = parser.parse_known_args()
-    cfgs = load_config_with_cli(args.cfg, args_list=cfg_args)  # skip --cfg
-
-    if cfgs.seed is not None:
-        if is_list(cfgs.seed):
-            assert len(cfgs.seed) == cfgs.num*cfgs.bs, 'seed list length should be equal to num*bs'
-            seeds = list(cfgs.seed)
-        else:
-            seeds = list(range(cfgs.seed, cfgs.seed+cfgs.num*cfgs.bs))
-    else:
-        seeds = [None]*(cfgs.num*cfgs.bs)
-
-    viser = VisualizerFast(cfgs)
-
-    for i in range(cfgs.num):
-        prompt = cfgs.prompt[i*cfgs.bs:(i+1)*cfgs.bs] if is_list(cfgs.prompt) else [cfgs.prompt]*cfgs.bs
-        negative_prompt = cfgs.neg_prompt[i*cfgs.bs:(i+1)*cfgs.bs] if is_list(cfgs.neg_prompt) else [cfgs.neg_prompt]*cfgs.bs
-        viser.vis_to_dir(prompt=prompt, negative_prompt=negative_prompt,
-                         seeds=seeds[i*cfgs.bs:(i+1)*cfgs.bs], save_cfg=cfgs.save.save_cfg, **cfgs.infer_args)