From c01c60e3c1a57028cc9c3992ac27120299e5f97f Mon Sep 17 00:00:00 2001 From: aszc-dev Date: Fri, 3 Nov 2023 01:05:28 +0100 Subject: [PATCH] Add progress bar and preview to LCM --- coreml_suite/lcm/lcm_sampler.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/coreml_suite/lcm/lcm_sampler.py b/coreml_suite/lcm/lcm_sampler.py index fe25bff..5b90321 100644 --- a/coreml_suite/lcm/lcm_sampler.py +++ b/coreml_suite/lcm/lcm_sampler.py @@ -3,6 +3,7 @@ import os import numpy as np import torch from diffusers.utils.torch_utils import randn_tensor +from tqdm import tqdm import comfy.utils import latent_preview @@ -137,14 +138,18 @@ class CoreMLSamplerLCM(CoreMLSampler): positive = positive[0][0] + callback = latent_preview.prepare_callback(patched_model, steps, None) torch.manual_seed(seed) - return self._sample(patched_model, steps, cfg, positive, latent_image, denoise) + return self._sample( + patched_model, steps, cfg, positive, latent_image, denoise, callback + ) - def _sample(self, model, steps, cfg, positive, latent_image, denoise): + def _sample( + self, model, steps, cfg, positive, latent_image, denoise, callback=None + ): device = get_torch_device() batch_size = latent_image["samples"].shape[0] - # callback = latent_preview.prepare_callback(model, steps, None) prompt_embeds = self.prepare_prompt_embeds(batch_size, positive) @@ -158,7 +163,8 @@ class CoreMLSamplerLCM(CoreMLSampler): ) # LCM MultiStep Sampling Loop: - for i, t in enumerate(timesteps): + iterator = tqdm(timesteps, desc="Core ML LCM Sampler", total=steps) + for i, t in enumerate(iterator): ts = torch.full((batch_size,), t, device=device, dtype=torch.float16) model_pred = model.model( @@ -173,6 +179,9 @@ class CoreMLSamplerLCM(CoreMLSampler): model_pred, i, t, latents, return_dict=False ) + if callback: + callback(i, denoised, latents, steps) + denoised = denoised.to(get_torch_device()) return ({"samples": denoised / 0.1825},)