coherent but fully denoised output
This commit is contained in:
RyanOnTheInside
+40
-68
@@ -3,6 +3,8 @@ import comfy.model_management
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
|
||||
class StreamCrossAttention:
|
||||
original_attention = None
|
||||
|
||||
@classmethod
|
||||
def get_attention_scores(cls, q, k, v, extra_options):
|
||||
"""
|
||||
@@ -11,55 +13,39 @@ class StreamCrossAttention:
|
||||
- KV cache preservation
|
||||
- Sliced attention fallback
|
||||
"""
|
||||
print(f"[DEBUG] Attention input q stats - min: {q.min():.4f}, max: {q.max():.4f}, mean: {q.mean():.4f}")
|
||||
print(f"[DEBUG] Attention input k stats - min: {k.min():.4f}, max: {k.max():.4f}, mean: {k.mean():.4f}")
|
||||
print(f"[DEBUG] Attention input v stats - min: {v.min():.4f}, max: {v.max():.4f}, mean: {v.mean():.4f}")
|
||||
|
||||
if not hasattr(cls, "attention_impl"):
|
||||
# Initialize optimized implementation once
|
||||
cls.attention_impl = cls.initialize_attention()
|
||||
# Get transformer options
|
||||
transformer_options = extra_options.get("transformer_options", {})
|
||||
|
||||
# Get cached key/values if available
|
||||
cache = extra_options.get("_stream_kv_cache", {})
|
||||
kv_cache = transformer_options.get("kv_cache", {})
|
||||
layer_id = extra_options["layer_id"]
|
||||
|
||||
if layer_id not in cache:
|
||||
cache[layer_id] = {"k": None, "v": None}
|
||||
|
||||
# Use cached KV if available
|
||||
if cache[layer_id]["k"] is not None:
|
||||
k = cache[layer_id]["k"]
|
||||
v = cache[layer_id]["v"]
|
||||
print(f"[DEBUG] Using cached k stats - min: {k.min():.4f}, max: {k.max():.4f}, mean: {k.mean():.4f}")
|
||||
print(f"[DEBUG] Using cached v stats - min: {v.min():.4f}, max: {v.max():.4f}, mean: {v.mean():.4f}")
|
||||
else: # First run or cache disabled
|
||||
cache[layer_id]["k"] = k
|
||||
cache[layer_id]["v"] = v
|
||||
# Check if we should use cached KV
|
||||
if kv_cache is not None and layer_id in kv_cache:
|
||||
cached = kv_cache[layer_id]
|
||||
k = cached["k"]
|
||||
v = cached["v"]
|
||||
else:
|
||||
# Store KV in cache if enabled
|
||||
if kv_cache is not None:
|
||||
kv_cache[layer_id] = {"k": k, "v": v}
|
||||
|
||||
# Use ComfyUI's memory-efficient attention as base
|
||||
output = optimized_attention(
|
||||
# Use ComfyUI's memory-efficient attention
|
||||
return optimized_attention(
|
||||
q, k, v,
|
||||
heads=extra_options["n_heads"],
|
||||
mask=extra_options["mask"],
|
||||
attn_precision=None, # Let ComfyUI handle precision
|
||||
skip_reshape=False, # Default reshape behavior
|
||||
skip_output_reshape=False # Default output reshape behavior
|
||||
mask=extra_options.get("mask"),
|
||||
)
|
||||
print(f"[DEBUG] Attention output stats - min: {output.min():.4f}, max: {output.max():.4f}, mean: {output.mean():.4f}")
|
||||
return output
|
||||
|
||||
@classmethod
|
||||
def initialize_attention(cls):
|
||||
"""Patch ComfyUI's attention modules with our optimized version"""
|
||||
from comfy.ldm.modules.attention import CrossAttention
|
||||
original_attention = CrossAttention.forward
|
||||
|
||||
if cls.original_attention is None:
|
||||
cls.original_attention = CrossAttention.forward
|
||||
|
||||
def stream_attention_forward(self, x, context=None, mask=None, **kwargs):
|
||||
# Convert to BHWC format if needed
|
||||
if x.shape[1] == self.dim_head * self.heads:
|
||||
x = x.view(x.shape[0], -1, self.dim_head).permute(0, 2, 1, 3)
|
||||
|
||||
# Use our optimized attention implementation
|
||||
q = self.to_q(x)
|
||||
context = context if context is not None else x
|
||||
k = self.to_k(context)
|
||||
@@ -69,8 +55,9 @@ class StreamCrossAttention:
|
||||
q, k, v,
|
||||
extra_options={
|
||||
"n_heads": self.heads,
|
||||
"layer_id": id(self), # Unique ID for KV caching
|
||||
"mask": mask
|
||||
"layer_id": id(self),
|
||||
"mask": mask,
|
||||
"transformer_options": kwargs.get("transformer_options", {})
|
||||
}
|
||||
)
|
||||
return self.to_out(out)
|
||||
@@ -90,7 +77,7 @@ class UNetBatchedStream:
|
||||
self.inner_model = model
|
||||
self.kv_cache = {}
|
||||
self.noise_base = None
|
||||
self.current_channels = None
|
||||
self.current_batch_size = None
|
||||
|
||||
def __call__(self, apply_model, args):
|
||||
"""
|
||||
@@ -106,46 +93,31 @@ class UNetBatchedStream:
|
||||
timestep = args["timestep"]
|
||||
c = args["c"]
|
||||
|
||||
print(f"[DEBUG] Initial x stats - min: {x.min():.4f}, max: {x.max():.4f}, mean: {x.mean():.4f}")
|
||||
# Get current batch size
|
||||
batch_size = x.shape[0]
|
||||
|
||||
# Get current channel count from input
|
||||
current_channels = x.shape[1]
|
||||
|
||||
# Reinitialize noise base if channels change
|
||||
if self.current_channels != current_channels:
|
||||
# Reinitialize noise base if batch size changes
|
||||
if self.current_batch_size != batch_size:
|
||||
self.noise_base = None
|
||||
self.current_channels = current_channels
|
||||
self.current_batch_size = batch_size
|
||||
|
||||
# Generate shared noise pattern matching current channels
|
||||
# Generate shared noise pattern for batch
|
||||
if self.noise_base is None:
|
||||
self.noise_base = torch.randn_like(x)
|
||||
print(f"[DEBUG] Generated noise_base - min: {self.noise_base.min():.4f}, max: {self.noise_base.max():.4f}, mean: {self.noise_base.mean():.4f}")
|
||||
|
||||
# Apply noise blending using ComfyUI's native sigma scheduling
|
||||
x = self.blend_noise(x, timestep)
|
||||
print(f"[DEBUG] After blend_noise - min: {x.min():.4f}, max: {x.max():.4f}, mean: {x.mean():.4f}")
|
||||
|
||||
# Pass through model
|
||||
# Get transformer options
|
||||
transformer_options = c.get("transformer_options", {})
|
||||
|
||||
# Add KV cache to transformer options
|
||||
transformer_options["kv_cache"] = self.kv_cache
|
||||
|
||||
# Pass through model with updated options
|
||||
return apply_model(
|
||||
x,
|
||||
timestep,
|
||||
c_concat=c.get("c_concat"),
|
||||
c_crossattn=c.get("c_crossattn"),
|
||||
control=c.get("control"),
|
||||
transformer_options=c.get("transformer_options", {}),
|
||||
cond_or_uncond=args.get("cond_or_uncond", None),
|
||||
kv_cache=self.kv_cache
|
||||
)
|
||||
def blend_noise(self, x, timesteps):
|
||||
"""ComfyUI-native noise blending"""
|
||||
# Convert timesteps to sigmas (this matches ComfyUI's internal scheduling)
|
||||
sigma = timesteps
|
||||
|
||||
# Simple sigma-based noise mixing
|
||||
noise_level = 1.0 / (1.0 + sigma)
|
||||
noise_level = noise_level.reshape(-1, 1, 1, 1)
|
||||
|
||||
# Mix clean latent with noise base
|
||||
x = x * (1.0 - noise_level) + self.noise_base * noise_level
|
||||
|
||||
return x
|
||||
transformer_options=transformer_options,
|
||||
cond_or_uncond=args.get("cond_or_uncond", None)
|
||||
)
|
||||
|
||||
+27
-52
@@ -71,11 +71,9 @@ class StreamDiffusionPipeline(ControlNodeBase):
|
||||
"latent_buffer": None, # Shape: [num_timesteps * frame_buffer_size, 4, H, W]
|
||||
"current_idx": 0,
|
||||
"t_indices": None,
|
||||
"alpha_prod_t_sqrt": None,
|
||||
"beta_prod_t_sqrt": None,
|
||||
"init_noise": None,
|
||||
"sigmas": None,
|
||||
"stock_noise": None, # For RCFG
|
||||
"timesteps": None,
|
||||
})
|
||||
|
||||
# Get latent tensor from dict
|
||||
@@ -93,71 +91,50 @@ class StreamDiffusionPipeline(ControlNodeBase):
|
||||
state["sigmas"] = sigmas
|
||||
|
||||
# Parse t_index_list and get timesteps
|
||||
state["t_indices"] = [int(x.strip()) for x in t_index_list.split(",")]
|
||||
t_indices = [int(x.strip()) for x in t_index_list.split(",")]
|
||||
|
||||
# Get actual timesteps from scheduler
|
||||
timesteps = []
|
||||
alpha_prod_t_sqrt_list = []
|
||||
beta_prod_t_sqrt_list = []
|
||||
for t_idx in t_indices:
|
||||
timesteps.append(len(sigmas) - 1 - t_idx)
|
||||
|
||||
# Convert sigmas to timesteps and calculate alpha/beta
|
||||
for t_idx in state["t_indices"]:
|
||||
sigma = sigmas[t_idx]
|
||||
# Convert sigma to timestep
|
||||
timestep = len(sigmas) - 1 - t_idx
|
||||
timesteps.append(timestep)
|
||||
|
||||
# Calculate alpha/beta from sigma
|
||||
alpha = 1.0 / (1.0 + sigma * sigma)
|
||||
beta = sigma * sigma / (1.0 + sigma * sigma)
|
||||
|
||||
alpha_prod_t_sqrt_list.append(alpha.sqrt())
|
||||
beta_prod_t_sqrt_list.append(beta.sqrt())
|
||||
|
||||
# Update t_indices to use actual timesteps
|
||||
state["t_indices"] = timesteps
|
||||
|
||||
# Stack and reshape for broadcasting
|
||||
state["alpha_prod_t_sqrt"] = torch.stack(alpha_prod_t_sqrt_list).view(-1, 1, 1, 1).to(device)
|
||||
state["beta_prod_t_sqrt"] = torch.stack(beta_prod_t_sqrt_list).view(-1, 1, 1, 1).to(device)
|
||||
# Store timesteps for later use
|
||||
state["timesteps"] = torch.tensor(timesteps, dtype=torch.long, device=device)
|
||||
state["t_indices"] = torch.tensor(t_indices, dtype=torch.long, device=device)
|
||||
|
||||
# Calculate total buffer size based on timesteps and frame buffer
|
||||
total_buffer_size = len(state["t_indices"]) * frame_buffer_size
|
||||
total_buffer_size = len(state["timesteps"]) * frame_buffer_size
|
||||
|
||||
# Initialize buffers if needed
|
||||
if state["latent_buffer"] is None:
|
||||
state["latent_buffer"] = torch.zeros((total_buffer_size, *samples.shape[1:]),
|
||||
dtype=samples.dtype,
|
||||
device=device)
|
||||
# Generate fixed noise pattern
|
||||
state["init_noise"] = torch.randn(samples.shape, dtype=samples.dtype, device=device)
|
||||
state["stock_noise"] = torch.zeros_like(samples) # For RCFG
|
||||
# Initialize stock noise for RCFG
|
||||
state["stock_noise"] = torch.zeros_like(samples)
|
||||
state["current_idx"] = 0
|
||||
|
||||
# Update buffer with current frame
|
||||
for i in range(len(state["t_indices"])):
|
||||
for i in range(len(state["timesteps"])):
|
||||
buffer_idx = i * frame_buffer_size + (state["current_idx"] % frame_buffer_size)
|
||||
if do_add_noise:
|
||||
# Add scaled noise using proper scheduler parameters
|
||||
alpha = state["alpha_prod_t_sqrt"][i]
|
||||
beta = state["beta_prod_t_sqrt"][i]
|
||||
state["latent_buffer"][buffer_idx] = (
|
||||
alpha * samples[0] + # Scale clean latent
|
||||
beta * state["init_noise"][0] # Use fixed noise pattern
|
||||
)
|
||||
# Get sigma for current timestep
|
||||
sigma = state["sigmas"][state["t_indices"][i]]
|
||||
# Add noise using sigma scheduling
|
||||
noise = torch.randn_like(samples)
|
||||
state["latent_buffer"][buffer_idx] = samples[0] + noise[0] * sigma
|
||||
else:
|
||||
state["latent_buffer"][buffer_idx] = samples[0]
|
||||
|
||||
# Create batched input for denoising
|
||||
batched = state["latent_buffer"].view(-1, *samples.shape[1:]) # [total_buffer_size, 4, H, W]
|
||||
timesteps = torch.tensor(state["t_indices"], dtype=torch.long, device=device)
|
||||
batched = state["latent_buffer"] # Already in right shape [total_buffer_size, 4, H, W]
|
||||
|
||||
# Get sigmas for selected timesteps
|
||||
timestep_sigmas = state["sigmas"][state["t_indices"]]
|
||||
|
||||
# Create sampler
|
||||
sampler = comfy.samplers.sampler_object(sampler_name)
|
||||
|
||||
# Get sigmas for selected timesteps
|
||||
timestep_sigmas = state["sigmas"][timesteps]
|
||||
|
||||
# Create custom CFG function for the sampler
|
||||
def custom_cfg(x, t, cond, uncond, cfg_scale, model_options={}):
|
||||
with torch.no_grad():
|
||||
@@ -172,13 +149,14 @@ class StreamDiffusionPipeline(ControlNodeBase):
|
||||
# Get unconditional prediction
|
||||
noise_pred_uncond = model.model.apply_model(x, t, uncond)
|
||||
elif cfg_type == "self" or cfg_type == "initialize":
|
||||
# Use stock noise for unconditional
|
||||
noise_pred_uncond = state["stock_noise"] * delta
|
||||
# Use stock noise for unconditional with delta scaling
|
||||
noise_pred_uncond = state["stock_noise"].repeat(x.shape[0], 1, 1, 1) * delta
|
||||
# Update stock noise for next iteration
|
||||
if cfg_type == "self":
|
||||
state["stock_noise"] = noise_pred_text.detach()
|
||||
elif cfg_type == "initialize" and t == timesteps[0]: # First step
|
||||
state["stock_noise"] = noise_pred_text.detach()
|
||||
# Update per batch element
|
||||
state["stock_noise"] = noise_pred_text.reshape(-1, frame_buffer_size, *noise_pred_text.shape[1:])[-1]
|
||||
elif cfg_type == "initialize" and t == timestep_sigmas[0]: # First step
|
||||
state["stock_noise"] = noise_pred_text[-1:] # Use last prediction
|
||||
|
||||
# Apply CFG
|
||||
return noise_pred_uncond + cfg_scale * (noise_pred_text - noise_pred_uncond)
|
||||
@@ -192,14 +170,11 @@ class StreamDiffusionPipeline(ControlNodeBase):
|
||||
if not hasattr(model, "_stream_wrapped"):
|
||||
model._stream_wrapped = UNetBatchedStream(model)
|
||||
model.set_model_unet_function_wrapper(model._stream_wrapped)
|
||||
|
||||
# Convert latent to BHWC format
|
||||
samples = samples.permute(0, 2, 3, 1) # BCHW -> BHWC
|
||||
|
||||
# Modified denoising call:
|
||||
denoised = comfy.samplers.sample(
|
||||
model,
|
||||
state["init_noise"].repeat(len(batched), 1, 1, 1),
|
||||
torch.randn_like(batched), # Fresh noise per batch
|
||||
positive,
|
||||
negative,
|
||||
guidance_scale,
|
||||
|
||||
Reference in New Issue
Block a user