File size: 7,352 Bytes
c65ecd1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
# -*- coding: utf-8 -*-
"""Collect calibration dataset."""
import os
from dataclasses import dataclass
import datasets
import torch
from omniconfig import configclass
from torch import nn
from tqdm import tqdm
from deepcompressor.app.diffusion.config import DiffusionPtqRunConfig
from deepcompressor.utils.common import hash_str_to_int, tree_map
from ...utils import get_control
from ..data import get_dataset
from .utils import CollectHook
def process(x: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
return torch.from_numpy(x.float().numpy()).to(dtype)
def collect(config: DiffusionPtqRunConfig, dataset: datasets.Dataset):
samples_dirpath = os.path.join(config.output.root, "samples")
caches_dirpath = os.path.join(config.output.root, "caches")
os.makedirs(samples_dirpath, exist_ok=True)
os.makedirs(caches_dirpath, exist_ok=True)
caches = []
pipeline = config.pipeline.build()
model = pipeline.unet if hasattr(pipeline, "unet") else pipeline.transformer
assert isinstance(model, nn.Module)
model.register_forward_hook(CollectHook(caches=caches), with_kwargs=True)
batch_size = config.eval.batch_size
print(f"In total {len(dataset)} samples")
print(f"Evaluating with batch size {batch_size}")
pipeline.set_progress_bar_config(desc="Sampling", leave=False, dynamic_ncols=True, position=1)
for batch in tqdm(
dataset.iter(batch_size=batch_size, drop_last_batch=False),
desc="Data",
leave=False,
dynamic_ncols=True,
total=(len(dataset) + batch_size - 1) // batch_size,
):
filenames = batch["filename"]
prompts = batch["prompt"]
seeds = [hash_str_to_int(name) for name in filenames]
generators = [torch.Generator(device=pipeline.device).manual_seed(seed) for seed in seeds]
pipeline_kwargs = config.eval.get_pipeline_kwargs()
task = config.pipeline.task
control_root = config.eval.control_root
if task in ["canny-to-image", "depth-to-image", "inpainting"]:
controls = get_control(
task,
batch["image"],
names=batch["filename"],
data_root=os.path.join(
control_root, collect_config.dataset_name, f"{dataset.config_name}-{config.eval.num_samples}"
),
)
if task == "inpainting":
pipeline_kwargs["image"] = controls[0]
pipeline_kwargs["mask_image"] = controls[1]
else:
pipeline_kwargs["control_image"] = controls
# Handle meta tensors by moving individual components
try:
pipeline = pipeline.to("cuda")
except NotImplementedError:
# Move individual pipeline components that have to_empty method
if hasattr(pipeline, 'transformer') and pipeline.transformer is not None:
try:
pipeline.transformer = pipeline.transformer.to("cuda")
except NotImplementedError:
pipeline.transformer = pipeline.transformer.to_empty(device="cuda")
if hasattr(pipeline, 'text_encoder') and pipeline.text_encoder is not None:
try:
pipeline.text_encoder = pipeline.text_encoder.to("cuda")
except NotImplementedError:
pipeline.text_encoder = pipeline.text_encoder.to_empty(device="cuda")
if hasattr(pipeline, 'text_encoder_2') and pipeline.text_encoder_2 is not None:
try:
pipeline.text_encoder_2 = pipeline.text_encoder_2.to("cuda")
except NotImplementedError:
pipeline.text_encoder_2 = pipeline.text_encoder_2.to_empty(device="cuda")
if hasattr(pipeline, 'vae') and pipeline.vae is not None:
try:
pipeline.vae = pipeline.vae.to("cuda")
except NotImplementedError:
pipeline.vae = pipeline.vae.to_empty(device="cuda")
result_images = pipeline(prompt=prompts, generator=generators, **pipeline_kwargs).images
num_guidances = (len(caches) // batch_size) // config.eval.num_steps
num_steps = len(caches) // (batch_size * num_guidances)
assert (
len(caches) == batch_size * num_steps * num_guidances
), f"Unexpected number of caches: {len(caches)} != {batch_size} * {config.eval.num_steps} * {num_guidances}"
for j, (filename, image) in enumerate(zip(filenames, result_images, strict=True)):
image.save(os.path.join(samples_dirpath, f"{filename}.png"))
for s in range(num_steps):
for g in range(num_guidances):
c = caches[s * batch_size * num_guidances + g * batch_size + j]
c["filename"] = filename
c["step"] = s
c["guidance"] = g
c = tree_map(lambda x: process(x), c)
torch.save(c, os.path.join(caches_dirpath, f"{filename}-{s:05d}-{g}.pt"))
caches.clear()
@configclass
@dataclass
class CollectConfig:
"""Configuration for collecting calibration dataset.
Args:
root (`str`, *optional*, defaults to `"datasets"`):
Root directory to save the collected dataset.
dataset_name (`str`, *optional*, defaults to `"qdiff"`):
Name of the collected dataset.
prompt_path (`str`, *optional*, defaults to `"prompts/qdiff.yaml"`):
Path to the prompt file.
num_samples (`int`, *optional*, defaults to `128`):
Number of samples to collect.
"""
root: str = "datasets"
dataset_name: str = "qdiff"
data_path: str = "prompts/qdiff.yaml"
num_samples: int = 128
if __name__ == "__main__":
parser = DiffusionPtqRunConfig.get_parser()
parser.add_config(CollectConfig, scope="collect", prefix="collect")
configs, _, unused_cfgs, unused_args, unknown_args = parser.parse_known_args()
ptq_config, collect_config = configs[""], configs["collect"]
assert isinstance(ptq_config, DiffusionPtqRunConfig)
assert isinstance(collect_config, CollectConfig)
if len(unused_cfgs) > 0:
print(f"Warning: unused configurations {unused_cfgs}")
if unused_args is not None:
print(f"Warning: unused arguments {unused_args}")
assert len(unknown_args) == 0, f"Unknown arguments: {unknown_args}"
collect_dirpath = os.path.join(
collect_config.root,
str(ptq_config.pipeline.dtype),
ptq_config.pipeline.name,
ptq_config.eval.protocol,
collect_config.dataset_name,
f"s{collect_config.num_samples}",
)
print(f"Saving caches to {collect_dirpath}")
dataset = get_dataset(
collect_config.data_path,
max_dataset_size=collect_config.num_samples,
return_gt=ptq_config.pipeline.task in ["canny-to-image"],
repeat=1,
)
ptq_config.output.root = collect_dirpath
os.makedirs(ptq_config.output.root, exist_ok=True)
collect(ptq_config, dataset=dataset)
|