|
|
|
|
|
"""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
|
|
|
|
|
|
|
|
|
try:
|
|
|
pipeline = pipeline.to("cuda")
|
|
|
except NotImplementedError:
|
|
|
|
|
|
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)
|
|
|
|
