lym00's picture
Upload calib.py
c65ecd1 verified
# -*- 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)