Upload calib.py

calib.py

@@ -0,0 +1,174 @@
+# -*- 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)