Remove CUDA specific stuff in previewer

Add maxed batch step mode for previews

Improve grid layout heuristics in previewer

Documentation improvements
This commit is contained in:
blepping
2024-01-28 07:35:28 -07:00
parent bef8ff691c
commit 59a4be8f57
8 changed files with 60 additions and 57 deletions
+1
View File
@@ -1,4 +1,5 @@
blehconfig.json
blehconfig.yaml
# Byte-compiled / optimized / DLL files
__pycache__/
+10 -9
View File
@@ -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:
+1 -1
View File
@@ -5,6 +5,6 @@
"max_batch": 4,
"max_batch_cols": 2,
"throttle_secs": 1,
"use_cuda": true
"maxed_batch_step_mode": false
}
}
+7
View File
@@ -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
+9
View File
@@ -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.
View File
+26 -43
View File
@@ -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
+6 -4
View File
@@ -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