diff --git a/lynx/nodes.py b/lynx/nodes.py index 725a9e3..c6ec994 100644 --- a/lynx/nodes.py +++ b/lynx/nodes.py @@ -198,12 +198,14 @@ class WanVideoAddLynxEmbeds: "embeds": ("WANVIDIMAGE_EMBEDS",), "ip_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the ip adapter face feature"}), "ref_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the reference feature"}), + "lynx_cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "If above 1.0 and main cfg_scale is above 1.0, run extra pass, default value 2.0"}), "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percent to apply the ref "}), "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percent to apply the ref "}), }, "optional": { + "vae": ("WANVAE", {"tooltip": "VAE model, only needed if ref_image is provided"}), "lynx_ip_embeds": ("LYNXIP", {"tooltip": "lynx face embeddings"}), - "ref_latent": ("LATENT",), + "ref_image": ("IMAGE",), "ref_text_embed": ("WANVIDEOTEXTEMBEDS",) } } @@ -213,16 +215,25 @@ class WanVideoAddLynxEmbeds: FUNCTION = "add" CATEGORY = "WanVideoWrapper" - def add(self, embeds, ip_scale, ref_scale, start_percent, end_percent, lynx_ip_embeds=None, ref_latent=None, ref_text_embed=None): - if ref_latent is not None and ref_text_embed is None: - raise ValueError("If ref_latent is provided, ref_text_embed must also be provided.") + def add(self, embeds, ip_scale, ref_scale, start_percent, end_percent, lynx_cfg_scale, vae=None, lynx_ip_embeds=None, ref_image=None, ref_text_embed=None): + if ref_image is not None and ref_text_embed is None: + raise ValueError("If ref_image is provided, ref_text_embed must also be provided.") + if ref_image is not None: + vae.to(device) + ref_image_in = (ref_image[..., :3].permute(3, 0, 1, 2) * 2 - 1).to(device, vae.dtype) + ref_latent = vae.encode([ref_image_in], device, tiled=False) + ref_latent_uncond = vae.encode([torch.zeros_like(ref_image_in)], device, tiled=False) + vae.to(offload_device) + new_entry = { "ip_x": lynx_ip_embeds["ip_x"] if lynx_ip_embeds is not None else None, "ip_x_uncond": lynx_ip_embeds["ip_x_uncond"] if lynx_ip_embeds is not None else None, - "ref_latent": ref_latent["samples"] if ref_latent is not None else None, + "ref_latent": ref_latent if ref_image is not None else None, + "ref_latent_uncond": ref_latent_uncond if ref_image is not None else None, "ref_text_embed": ref_text_embed if ref_text_embed is not None else None, "ip_scale": ip_scale, "ref_scale": ref_scale, + "cfg_scale": lynx_cfg_scale, "start_percent": start_percent, "end_percent": end_percent, } diff --git a/nodes_sampler.py b/nodes_sampler.py index 11fec5c..b785382 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -1000,10 +1000,15 @@ class WanVideoSampler: lynx_embeds = lynx_embeds.copy() log.info("Using Lynx embeddings", lynx_embeds) lynx_ref_latent = lynx_embeds.get("ref_latent", None) + lynx_ref_latent_uncond = lynx_embeds.get("ref_latent_uncond", None) lynx_ref_text_embed = lynx_embeds.get("ref_text_embed", None) + lynx_cfg_scale = lynx_embeds.get("cfg_scale", 1.0) + if not isinstance(lynx_cfg_scale, list): + lynx_cfg_scale = [lynx_cfg_scale] * (steps + 1) if lynx_ref_latent is not None: lynx_ref_latent = lynx_ref_latent[0] + lynx_ref_latent_uncond = lynx_ref_latent_uncond[0] lynx_embeds["ref_feature_extractor"] = True log.info(f"Lynx ref latent shape: {lynx_ref_latent.shape}") log.info("Extracting Lynx ref buffer...") @@ -1016,7 +1021,7 @@ class WanVideoSampler: ) if not math.isclose(cfg[0], 1.0): lynx_ref_buffer_uncond = transformer( - [lynx_ref_latent.to(device, dtype) * 0], + [lynx_ref_latent_uncond.to(device, dtype)], torch.tensor([0], device=device), lynx_ref_text_embed["prompt_embeds"], seq_len=math.ceil((lynx_ref_latent.shape[2] * lynx_ref_latent.shape[3]) / 4 * lynx_ref_latent.shape[1]), @@ -1349,6 +1354,8 @@ class WanVideoSampler: #phantom if use_phantom and not math.isclose(phantom_cfg_scale[idx], 1.0): + if cache_state is not None and len(cache_state) != 3: + cache_state.append(None) noise_pred_phantom, cache_state_phantom = transformer( context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None, **base_params) @@ -1369,9 +1376,9 @@ class WanVideoSampler: else: # multitalk base_params['multitalk_audio'] = torch.zeros_like(multitalk_audio_input)[-1:] audio_context = negative_embeds - + base_params['is_uncond'] = False noise_pred_no_audio, cache_state_audio = transformer( - context=audio_context, is_uncond=False, + context=audio_context, pred_id=cache_state[2] if cache_state else None, vace_data=vace_data, **base_params) @@ -1380,6 +1387,18 @@ class WanVideoSampler: + cfg_scale * (noise_pred_no_audio[0] - noise_pred_uncond) + audio_cfg_scale[idx] * (noise_pred_cond - noise_pred_no_audio[0])) return noise_pred, [cache_state_cond, cache_state_uncond, cache_state_audio] + #lynx + if lynx_embeds is not None and not math.isclose(lynx_cfg_scale[idx], 1.0): + base_params['is_uncond'] = False + if cache_state is not None and len(cache_state) != 3: + cache_state.append(None) + noise_pred_lynx, cache_state_lynx = transformer( + context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None, + **base_params) + + noise_pred = (noise_pred_uncond + lynx_cfg_scale[idx] * (noise_pred_lynx[0] - noise_pred_uncond) + + cfg_scale * (noise_pred_cond - noise_pred_lynx[0])) + return noise_pred, [cache_state_cond, cache_state_uncond, cache_state_lynx] #batched else: