Update README.md

README.md

@@ -293,7 +293,98 @@ Potential fix: app.diffusion.nn.struct.py
 Potential Fix: app.diffusion.dataset.collect.calib.py
 
 ```python
+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()
 ```
 
 References