diff --git a/.gitignore b/.gitignore index ce918b2..aabefd5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ blehconfig.json +blehconfig.yaml # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index e774ebf..fcefa6b 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # BLEH -ComfyUI nodes collection... eventually. +A ComfyUI nodes collection... eventually. ## Features @@ -9,18 +9,17 @@ ComfyUI nodes collection... eventually. ## Configuration -Copy `blehconfig.json.example` to `blehconfig.json` in the node repo directory and edit the copy. +Copy either `blehconfig.yaml.example` or `blehconfig.json.example` to `blehconfig.yaml` or `blehconfig.json` respectively and edit the copy. When loading configuration, the YAML file will be prioritized if it exists and Python has YAML support. -Restart ComfyUI to apply new changes. +Restart ComfyUI to apply any new changes. ### Better TAESD previews * Supports setting max preview size (ComfyUI default is hardcoded to 512 max). * Supports showing previews for more than the first latent in the batch. * Supports throttling previews. Do you really need your expensive TAESD preview to get updated 3 times a second? -* Supports using CUDA streams to avoid waiting for a synchronize. Increases speed slightly at the cost of higher VRAM usage. For comparison, a batch of 8 768x768 images with throttle at `0.5` sec is `1.29s/it` with it on and `1.42s/it` with it off for me. -Current defaults from `blehconfig.json` +Current defaults: |Key|Default|Description| |-|-|-| @@ -28,17 +27,19 @@ Current defaults from `blehconfig.json` |`max_size`|`768`|Max width or height for previews. Note this does not affect TAESD decoding, just the preview image| |`max_batch`|`4`|Max number of latents in a batch to preview| |`max_batch_cols`|`2`|Max number of columns to use when previewing batches| -|`throttle_secs`|`1`|Max frequency to decode the latents for previewing. `0.25` would be every 1/4 sec, `2` would be only once every two seconds| -|`use_cuda`|`false`|Use special logic for CUDA (and maybe pretend-CUDA like ROCM) to reduce the performance impact of preview generation| +|`throttle_secs`|`2`|Max frequency to decode the latents for previewing. `0.25` would be every quarter second, `2` would be once every two seconds| +|`maxed_batch_step_mode`|`false`|When `false`, you will see the first `max_batch` previews, when `true` you will see previews spread across the batch| -I would recommend setting `throttle_secs` to something relatively high like 5-10 sec especially if you are generating batches at high resolution. `use_cuda` now defaults to `false` as it can substantially increase VRAM requirements. +These defaults are conservative. I would recommend setting `throttle_secs` to something relatively high (like 5-10) especially if you are generating batches at high resolution. + +Slightly more detailed explanation for `maxed_batch_step_mode`: If max previews is set to `3` and the batch size is `15` you will see previews for indexes `0, 5, 10`. Or to put it a different way, it steps through the batch by `batch_size / max_previews` rounded up. This behavior may be useful for previewing generations with a high batch count like when using AnimateDiff. ### BlehHyperTile Adds the ability to set a seed and timestep range that HyperTile gets applied for. *Not* well tested, and I just assumed the Inspire version works which may or may not be the case. **Note**: Timesteps start from 999 and count down to 0 and also are not necessarily linear. Exactly what sampling step a timestep applies -to is left as an exercise for you, dear node user. +to is left as an exercise for you, dear node user. As an example, Karras and exponentially samplers essentially rush to low timesteps and spend quite a bit of time there. HyperTile credits: diff --git a/blehconfig.json.example b/blehconfig.json.example index 0076195..2af1131 100644 --- a/blehconfig.json.example +++ b/blehconfig.json.example @@ -5,6 +5,6 @@ "max_batch": 4, "max_batch_cols": 2, "throttle_secs": 1, - "use_cuda": true + "maxed_batch_step_mode": false } } diff --git a/blehconfig.yaml.example b/blehconfig.yaml.example new file mode 100644 index 0000000..3c5c18a --- /dev/null +++ b/blehconfig.yaml.example @@ -0,0 +1,7 @@ +betterTaesdPreviews: + enabled: true + max_size: 768 + max_batch: 4 + max_batch_cols: 2 + throttle_secs: 1 + maxed_batch_step_mode: false diff --git a/changelog.md b/changelog.md new file mode 100644 index 0000000..0ff9515 --- /dev/null +++ b/changelog.md @@ -0,0 +1,9 @@ +# Changes + +Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top. + +## 20240128 + +* Removed CUDA-specific stuff from TAESD previewer as the performance gains were marginal and it had a major effect on VRAM usage. +* (Hopefully) improved heuristics for batch preview layout. +* Added `maxed_batch_step_mode` setting for TAESD previewer. diff --git a/py/__init__.py b/py/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/py/betterTaesdPreview.py b/py/betterTaesdPreview.py index 95a45f7..e8680f7 100644 --- a/py/betterTaesdPreview.py +++ b/py/betterTaesdPreview.py @@ -20,11 +20,6 @@ class BetterTAESDPreviewer(_ORIG_PREVIEWER): self.stream = None self.prev_work = None self.cpudev = torch.device("cpu") - self.use_cuda = ( - SETTINGS.btp_use_cuda - and hasattr(torch, "cuda") - and torch.cuda.is_available() - ) def decode_latent_to_preview_image(self, preview_format, x0): preview_image = self.decode_latent_to_preview(x0) @@ -44,49 +39,35 @@ class BetterTAESDPreviewer(_ORIG_PREVIEWER): return False def _decode_latent(self, x0): - samples = (self.taesd.decode(x0[: SETTINGS.btp_max_batch]) + 1.0) / 2.0 + max_batch = SETTINGS.btp_max_batch + batch_size = x0.shape[0] + if not SETTINGS.btp_maxed_batch_step_mode: + indexes = range(min(max_batch, batch_size)) + else: + indexes = range( + 0, + batch_size, + math.ceil(batch_size / max_batch), + )[:max_batch] + samples = (self.taesd.decode(x0[indexes, :]) + 1.0) / 2.0 samples = torch.clamp(samples, min=0.0, max=1.0) * 255.0 return samples.to(dtype=torch.uint8).detach() - def decode_latent_to_preview(self, x0): - use_cached = self.check_use_cached() - if x0.device == self.cpudev or not self.use_cuda: - return ( - self.cached if use_cached else self._decode_latent_to_preview_nocuda(x0) - ) - if self.stream is None: - self.stream = torch.cuda.Stream() - elif not self.stream.query(): - return self.cached or self.blank - work = None - if self.prev_work is not None: - # We will only arrive here if the stream is ready. Sync just to be safe, should be instant. - self.stream.synchronize() - work = self.prev_work - del self.prev_work - result = self.work_to_image(work) if work is not None else self.blank - if use_cached: - return result - # The original stream may be still processing the current step. - orig_stream = torch.cuda.current_stream() - self.stream.wait_stream(orig_stream) - try: - torch.cuda.set_stream(self.stream) - self.prev_work = self._decode_latent(x0).to( - device=self.cpudev, - non_blocking=True, - ) - finally: - torch.cuda.set_stream(orig_stream) - return result - def calc_cols_rows(self, batch_size, width, height): - ratio = width / height - cols = min(math.ceil(batch_size / 2), SETTINGS.btp_max_batch_cols) + max_cols = SETTINGS.btp_max_batch_cols + ratio = height / width + if ratio >= 1.75: + # Very tall images - prioritize horizontal layout. + cols = min(batch_size, max_cols) + elif ratio <= 0.5: + # Very wide images - prioritize vertical layout. + cols = min(math.ceil(batch_size / 4), max_cols) + else: + cols = min(math.ceil(batch_size / ratio), max_cols) rows = math.ceil(batch_size / cols) return cols, rows - def work_to_image(self, samples): + def decoded_to_image(self, samples): samples = tuple(np.moveaxis(x, 0, 2) for x in samples.numpy()) batch_size = len(samples) height, width, _ = samples[0].shape @@ -104,8 +85,10 @@ class BetterTAESDPreviewer(_ORIG_PREVIEWER): self.cached = result return result - def _decode_latent_to_preview_nocuda(self, x0): - return self.work_to_image(self._decode_latent(x0).cpu()) + def decode_latent_to_preview(self, x0): + if self.check_use_cached(): + return self.cached + return self.decoded_to_image(self._decode_latent(x0).cpu()) latent_preview.TAESDPreviewerImpl = BetterTAESDPreviewer diff --git a/py/settings.py b/py/settings.py index 10793d0..43f76ad 100644 --- a/py/settings.py +++ b/py/settings.py @@ -11,11 +11,11 @@ class Settings: self.btp_enabled = False else: self.btp_enabled = True - self.btp_max_size = btp.get("max_size", 768) - self.btp_max_batch = btp.get("max_batch", 4) - self.btp_max_batch_cols = btp.get("max_batch_cols", 2) + self.btp_max_size = max(8, btp.get("max_size", 768)) + self.btp_max_batch = max(1, btp.get("max_batch", 4)) + self.btp_max_batch_cols = max(1, btp.get("max_batch_cols", 2)) self.btp_throttle_secs = btp.get("throttle_secs", 1) - self.btp_use_cuda = btp.get("use_cuda", True) + self.btp_maxed_batch_step_mode = btp.get("maxed_batch_step_mode", False) def get_cfg_path(self, filename): my_path = Path.resolve(Path(__file__).parent) @@ -27,6 +27,7 @@ class Settings: try: with Path.open(self.get_cfg_path(filename)) as fp: self.update(json.load(fp)) + return True except OSError: return False @@ -36,6 +37,7 @@ class Settings: with Path.open(self.get_cfg_path(filename)) as fp: self.update(yaml.safe_load(fp)) + return True except (OSError, ImportError): return False