Support loading AnyFlow embedder and weights as PEFT LoRA

demo.py

@@ -27,6 +27,7 @@
 from far.models.transformer_far_wan_model import FAR_Wan_Transformer3DModel
 from far.pipelines.pipeline_far_wan_anyflow import FARWanAnyFlowPipeline
 from far.pipelines.pipeline_wan_anyflow import WanAnyFlowPipeline
+from far.utils.lora_adapter import load_transformer_lora
 from far.utils.video_util import select_frame_indices
 from far.utils.vis_util import draw_rectangle
 
@@ -43,10 +44,15 @@ class DemoConfig:
     task_type: str = 't2v'
     # Where to write demo_*.mp4.
     save_dir: str = MISSING
+    # Optional PEFT LoRA safetensors path. AnyFlow sidecar tensors are loaded when present.
+    lora_path: str | None = None
+    lora_adapter_name: str = 'anyflow'
 
 
-def inference_causal_demo(model_path, task_type, save_dir):
+def inference_causal_demo(model_path, task_type, save_dir, lora_path=None, lora_adapter_name='anyflow'):
     transformer = FAR_Wan_Transformer3DModel.from_pretrained(model_path, subfolder='transformer')
+    if lora_path:
+        load_transformer_lora(transformer, lora_path, adapter_name=lora_adapter_name)
     pipeline = FARWanAnyFlowPipeline.from_pretrained(model_path, transformer=transformer).to('cuda', dtype=torch.bfloat16)
 
     os.makedirs(save_dir, exist_ok=True)
@@ -110,8 +116,10 @@ def inference_causal_demo(model_path, task_type, save_dir):
         raise NotImplementedError
 
 
-def inference_bidirectional_demo(model_path, task_type, save_dir):
+def inference_bidirectional_demo(model_path, task_type, save_dir, lora_path=None, lora_adapter_name='anyflow'):
     transformer = FAR_Wan_Transformer3DModel.from_pretrained(model_path, subfolder='transformer')
+    if lora_path:
+        load_transformer_lora(transformer, lora_path, adapter_name=lora_adapter_name)
     pipeline = WanAnyFlowPipeline.from_pretrained(model_path, transformer=transformer).to('cuda', dtype=torch.bfloat16)
     os.makedirs(save_dir, exist_ok=True)
 
@@ -137,9 +145,21 @@ def inference_bidirectional_demo(model_path, task_type, save_dir):
         OmegaConf.from_cli(),
     )
     if 'AnyFlow-FAR' in cfg.model_path:
-        inference_causal_demo(cfg.model_path, task_type=cfg.task_type, save_dir=cfg.save_dir)
+        inference_causal_demo(
+            cfg.model_path,
+            task_type=cfg.task_type,
+            save_dir=cfg.save_dir,
+            lora_path=cfg.lora_path,
+            lora_adapter_name=cfg.lora_adapter_name,
+        )
     elif 'AnyFlow-Wan' in cfg.model_path:
-        inference_bidirectional_demo(cfg.model_path, task_type=cfg.task_type, save_dir=cfg.save_dir)
+        inference_bidirectional_demo(
+            cfg.model_path,
+            task_type=cfg.task_type,
+            save_dir=cfg.save_dir,
+            lora_path=cfg.lora_path,
+            lora_adapter_name=cfg.lora_adapter_name,
+        )
     else:
         raise NotImplementedError
 
@@ -150,6 +170,7 @@ def inference_bidirectional_demo(model_path, task_type, save_dir):
   model_path   — Diffusers folder path.
   task_type    — t2v | ti2v | tv2v (default: t2v).
   save_dir     — Output directory for demo_*.mp4.
+  lora_path    — Optional PEFT LoRA safetensors path.
 
 Example (from repository root):
 python demo.py \

far/utils/lora_adapter.py

@@ -0,0 +1,111 @@
+# Copyright 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from pathlib import Path
+from typing import Optional
+
+import torch
+from safetensors import safe_open
+
+
+ANYFLOW_SIDECAR_PREFIXES = ("condition_embedder.delta_embedder.",)
+
+
+def enable_anyflow_time_conditioning(transformer, *, gate_value: float = 0.25, deltatime_type: str = "r") -> None:
+    if deltatime_type not in {"r", "t-r"}:
+        raise ValueError("AnyFlow deltatime_type must be 'r' or 't-r'.")
+    condition_embedder = getattr(transformer, "condition_embedder", None)
+    if condition_embedder is not None and hasattr(condition_embedder, "delta_embedder"):
+        return
+    if not hasattr(transformer, "setup_flowmap_model"):
+        raise ValueError("Transformer does not support AnyFlow time conditioning.")
+    transformer.register_to_config(gate_value=float(gate_value), deltatime_type=deltatime_type)
+    transformer.setup_flowmap_model()
+
+
+def _copy_sidecar_tensor(transformer, key: str, tensor: torch.Tensor, *, prefix: str) -> bool:
+    sidecar_prefix = f"{prefix}."
+    if not key.startswith(sidecar_prefix):
+        return False
+    model_key = key.removeprefix(sidecar_prefix)
+    if not model_key.startswith(ANYFLOW_SIDECAR_PREFIXES):
+        return False
+
+    model_state = transformer.state_dict()
+    try:
+        destination = model_state[model_key]
+    except KeyError as exc:
+        raise ValueError(
+            f"LoRA file contains AnyFlow sidecar tensor `{key}`, but the transformer does not have `{model_key}`. "
+            "Call enable_anyflow_time_conditioning(...) before loading this adapter."
+        ) from exc
+    if destination.shape != tensor.shape:
+        raise ValueError(
+            f"Shape mismatch for AnyFlow sidecar tensor `{key}`: "
+            f"model {tuple(destination.shape)} vs file {tuple(tensor.shape)}."
+        )
+    destination.copy_(tensor.to(device=destination.device, dtype=destination.dtype))
+    return True
+
+
+def load_transformer_lora(
+    transformer,
+    lora_path: str | Path,
+    *,
+    adapter_name: str = "anyflow",
+    prefix: str = "transformer",
+    gate_value: Optional[float] = None,
+    deltatime_type: Optional[str] = None,
+) -> None:
+    lora_path = Path(lora_path).expanduser()
+    lora_state = OrderedDict()
+
+    with safe_open(lora_path, framework="pt", device="cpu") as handle:
+        metadata = handle.metadata() or {}
+        sidecar_gate = float(gate_value if gate_value is not None else metadata.get("anyflow_gate_value", 0.25))
+        sidecar_deltatime = str(deltatime_type or metadata.get("anyflow_deltatime_type", "r"))
+        has_sidecar = metadata.get("simpletuner_anyflow_sidecar") == "true" or any(
+            "condition_embedder.delta_embedder." in key for key in handle.keys()
+        )
+        if has_sidecar:
+            enable_anyflow_time_conditioning(
+                transformer,
+                gate_value=sidecar_gate,
+                deltatime_type=sidecar_deltatime,
+            )
+
+        for key in handle.keys():
+            tensor = handle.get_tensor(key)
+            if has_sidecar and _copy_sidecar_tensor(transformer, key, tensor, prefix=prefix):
+                continue
+            if ".lora_A." in key or ".lora_B." in key or ".alpha" in key or ".lora_alpha" in key:
+                lora_state[key] = tensor
+
+    ordered_lora_state = OrderedDict()
+    for key, tensor in lora_state.items():
+        if ".lora_A." in key or ".lora_B." in key:
+            ordered_lora_state[key] = tensor
+    for key, tensor in lora_state.items():
+        if key not in ordered_lora_state:
+            ordered_lora_state[key] = tensor
+
+    if not ordered_lora_state:
+        raise ValueError(f"No PEFT LoRA tensors found in {lora_path}.")
+    transformer.load_lora_adapter(ordered_lora_state, prefix=prefix, adapter_name=adapter_name)
+    transformer.set_adapter(adapter_name)