diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index ef6d940..baf2d5f 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -11,6 +11,11 @@ from safetensors import safe_open from torch import nn from torch.nn import functional as F +try: + import comfy_kitchen # noqa: F401 +except Exception: + comfy_kitchen = None + IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) @@ -19,6 +24,10 @@ LATENTS_STD = (1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.75 def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor: + if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}: + if weight is None: + return F.rms_norm(x, (x.shape[-1],), eps=eps) + return F.rms_norm(x, weight.shape, weight=weight.to(device=x.device, dtype=x.dtype), eps=eps) result = x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps).to(x.dtype) return result if weight is None else result * weight.to(dtype=x.dtype) @@ -72,6 +81,11 @@ class FeedForward(nn.Module): def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor: """Apply the reference split-half RoPE layout to leading rotary channels.""" + if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}: + try: + return torch.ops.comfy_kitchen.apply_rope_split_half1(x, table) + except Exception: + pass pairs = table.shape[-3] rot = pairs * 2 first, second = x[..., :pairs], x[..., pairs:rot]