nodes/city96-ComfyUI_DiT
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Live preview + progress bar

Browse Source
This commit is contained in:
City
2023-09-05 19:06:44 +02:00
parent 03f06cc801
commit d01bf62c13
2 changed files with 15 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+4 -1
View File
@@ -6,6 +6,7 @@ import comfy.model_management
import comfy.model_patcher
import comfy.utils
import comfy.latent_formats
import latent_preview
from .models import DiT_models
from .diffusion import create_diffusion
@@ -121,6 +122,8 @@ class DiTSampler:
# pre
comfy.model_management.load_model_gpu(model)
real_model = model.model
pbar = comfy.utils.ProgressBar(steps)
previewer = latent_preview.get_previewer(device, model.model.latent_format)
# Create sampling noise:
z = torch.randn(batch_size, 4, real_model.latent_size, real_model.latent_size, device=device)
@@ -134,7 +137,7 @@ class DiTSampler:
# Sample images:
samples = diffusion.p_sample_loop(
model.model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
model.model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, pbar=pbar, previewer=previewer, device=device
)
samples, _ = samples.chunk(2, dim=0) # Remove null class samples
samples = real_model.latent_format.process_out(samples.to(torch.float32))