# -*- 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)