diffsynth-engine 0.6.1.dev9__tar.gz → 0.6.1.dev11__tar.gz

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.1.dev9
+Version: 0.6.1.dev11
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/__init__.py

@@ -14,6 +14,9 @@ from .configs import (
     AttnImpl,
     ControlNetParams,
     ControlType,
+    QwenImageControlNetParams,
+    QwenImageControlType,
+    LoraConfig,
 )
 from .pipelines import (
     SDImagePipeline,
@@ -58,6 +61,8 @@ __all__ = [
     "AttnImpl",
     "ControlNetParams",
     "ControlType",
+    "QwenImageControlNetParams",
+    "QwenImageControlType",
     "SDImagePipeline",
     "SDControlNet",
     "SDXLImagePipeline",
@@ -74,6 +79,7 @@ __all__ = [
     "FluxIPAdapterRefTool",
     "FluxReplaceByControlTool",
     "FluxReduxRefTool",
+    "LoraConfig",
     "fetch_model",
     "fetch_modelscope_model",
     "register_fetch_modelscope_model",

diffsynth_engine-0.6.1.dev11/diffsynth_engine/algorithm/noise_scheduler/base_scheduler.py

@@ -0,0 +1,29 @@
+import torch
+
+
+def append_zero(x):
+    return torch.cat([x, x.new_zeros([1])])
+
+
+class BaseScheduler:
+    def __init__(self):
+        self._stored_config = {}
+
+    def store_config(self):
+        self._stored_config = {
+            config_name: config_value
+            for config_name, config_value in vars(self).items()
+            if not config_name.startswith("_")
+        }
+
+    def update_config(self, config_dict):
+        for config_name, new_value in config_dict.items():
+            if hasattr(self, config_name):
+                setattr(self, config_name, new_value)
+
+    def restore_config(self):
+        for config_name, config_value in self._stored_config.items():
+            setattr(self, config_name, config_value)
+
+    def schedule(self, num_inference_steps: int):
+        raise NotImplementedError()

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/algorithm/noise_scheduler/flow_match/recifited_flow.py

@@ -12,16 +12,23 @@ class RecifitedFlowScheduler(BaseScheduler):
     def __init__(
         self,
         shift=1.0,
-        sigma_min=0.001,
-        sigma_max=1.0,
+        sigma_min=None,
+        sigma_max=None,
         num_train_timesteps=1000,
         use_dynamic_shifting=False,
+        shift_terminal=None,
+        exponential_shift_mu=None,
     ):
+        super().__init__()
         self.shift = shift
         self.sigma_min = sigma_min
         self.sigma_max = sigma_max
         self.num_train_timesteps = num_train_timesteps
         self.use_dynamic_shifting = use_dynamic_shifting
+        self.shift_terminal = shift_terminal
+        # static mu for distill model
+        self.exponential_shift_mu = exponential_shift_mu
+        self.store_config()
 
     def _sigma_to_t(self, sigma):
         return sigma * self.num_train_timesteps
@@ -35,21 +42,30 @@ class RecifitedFlowScheduler(BaseScheduler):
     def _shift_sigma(self, sigma: torch.Tensor, shift: float):
         return shift * sigma / (1 + (shift - 1) * sigma)
 
+    def _stretch_shift_to_terminal(self, sigma: torch.Tensor):
+        one_minus_z = 1 - sigma
+        scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+        return 1 - (one_minus_z / scale_factor)
+
     def schedule(
         self,
         num_inference_steps: int,
         mu: float | None = None,
-        sigma_min: float | None = None,
-        sigma_max: float | None = None,
+        sigma_min: float = 0.001,
+        sigma_max: float = 1.0,
         append_value: float = 0,
     ):
-        sigma_min = self.sigma_min if sigma_min is None else sigma_min
-        sigma_max = self.sigma_max if sigma_max is None else sigma_max
+        sigma_min = sigma_min if self.sigma_min is None else self.sigma_min
+        sigma_max = sigma_max if self.sigma_max is None else self.sigma_max
         sigmas = torch.linspace(sigma_max, sigma_min, num_inference_steps)
+        if self.exponential_shift_mu is not None:
+            mu = self.exponential_shift_mu
         if self.use_dynamic_shifting:
             sigmas = self._time_shift(mu, 1.0, sigmas)  # FLUX
         else:
             sigmas = self._shift_sigma(sigmas, self.shift)
+        if self.shift_terminal is not None:
+            sigmas = self._stretch_shift_to_terminal(sigmas)
         timesteps = sigmas * self.num_train_timesteps
         sigmas = append(sigmas, append_value)
         return sigmas, timesteps

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/configs/__init__.py

@@ -17,9 +17,15 @@ from .pipeline import (
     WanStateDicts,
     WanS2VStateDicts,
     QwenImageStateDicts,
+    LoraConfig,
     AttnImpl,
 )
-from .controlnet import ControlType, ControlNetParams
+from .controlnet import (
+    ControlType,
+    ControlNetParams,
+    QwenImageControlNetParams,
+    QwenImageControlType,
+)
 
 __all__ = [
     "BaseConfig",
@@ -40,7 +46,10 @@ __all__ = [
     "WanStateDicts",
     "WanS2VStateDicts",
     "QwenImageStateDicts",
-    "AttnImpl",
+    "QwenImageControlType",
+    "QwenImageControlNetParams",
     "ControlType",
     "ControlNetParams",
+    "LoraConfig",
+    "AttnImpl",
 ]

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/configs/controlnet.py

@@ -34,3 +34,16 @@ class ControlNetParams:
     control_start: float = 0
     control_end: float = 1
     processor_name: Optional[str] = None  # only used for sdxl controlnet union now
+
+
+class QwenImageControlType(Enum):
+    eligen = "eligen"
+    in_context = "in_context"
+
+
+@dataclass
+class QwenImageControlNetParams:
+    image: ImageType
+    model: str
+    control_type: QwenImageControlType
+    scale: float = 1.0

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/configs/pipeline.py

@@ -365,3 +365,9 @@ def init_parallel_config(config: FluxPipelineConfig | QwenImagePipelineConfig |
         config.tp_degree = 1
     else:
         raise ValueError("sp_ulysses_degree and sp_ring_degree must be specified together")
+
+
+@dataclass
+class LoraConfig:
+    scale: float
+    scheduler_config: Optional[Dict] = None

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/models/hunyuan3d/dino_image_encoder.py

@@ -2,7 +2,7 @@ import torch.nn as nn
 import torchvision.transforms as transforms
 import collections.abc
 import math
-from typing import Optional, Tuple, Dict
+from typing import Optional, Dict
 
 import torch
 from diffsynth_engine.models.base import PreTrainedModel, StateDictConverter
@@ -112,7 +112,9 @@ class Dinov2SelfAttention(nn.Module):
     def __init__(self, hidden_size: int, num_attention_heads: int, qkv_bias: bool) -> None:
         super().__init__()
         if hidden_size % num_attention_heads != 0:
-            raise ValueError(f"hidden_size {hidden_size} is not a multiple of num_attention_heads {num_attention_heads}.")
+            raise ValueError(
+                f"hidden_size {hidden_size} is not a multiple of num_attention_heads {num_attention_heads}."
+            )
 
         self.num_attention_heads = num_attention_heads
         self.attention_head_size = int(hidden_size / num_attention_heads)

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/models/qwen_image/qwen_image_dit.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn as nn
-from typing import Any, Dict, Tuple, Union, Optional
+from typing import Any, Dict, List, Tuple, Union, Optional
 from einops import rearrange
 
 from diffsynth_engine.models.base import StateDictConverter, PreTrainedModel
@@ -190,7 +190,8 @@ class QwenDoubleStreamAttention(nn.Module):
         self,
         image: torch.FloatTensor,
         text: torch.FloatTensor,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        attn_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
         img_q, img_k, img_v = self.to_q(image), self.to_k(image), self.to_v(image)
         txt_q, txt_k, txt_v = self.add_q_proj(text), self.add_k_proj(text), self.add_v_proj(text)
@@ -206,8 +207,8 @@ class QwenDoubleStreamAttention(nn.Module):
         img_q, img_k = self.norm_q(img_q), self.norm_k(img_k)
         txt_q, txt_k = self.norm_added_q(txt_q), self.norm_added_k(txt_k)
 
-        if image_rotary_emb is not None:
-            img_freqs, txt_freqs = image_rotary_emb
+        if rotary_emb is not None:
+            img_freqs, txt_freqs = rotary_emb
             img_q = apply_rotary_emb_qwen(img_q, img_freqs)
             img_k = apply_rotary_emb_qwen(img_k, img_freqs)
             txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs)
@@ -221,7 +222,7 @@ class QwenDoubleStreamAttention(nn.Module):
         joint_k = joint_k.transpose(1, 2)
         joint_v = joint_v.transpose(1, 2)
 
-        joint_attn_out = attention_ops.attention(joint_q, joint_k, joint_v, **self.attn_kwargs)
+        joint_attn_out = attention_ops.attention(joint_q, joint_k, joint_v, attn_mask=attn_mask, **self.attn_kwargs)
 
         joint_attn_out = rearrange(joint_attn_out, "b s h d -> b s (h d)").to(joint_q.dtype)
 
@@ -285,7 +286,8 @@ class QwenImageTransformerBlock(nn.Module):
         image: torch.Tensor,
         text: torch.Tensor,
         temb: torch.Tensor,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        attn_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         img_mod_attn, img_mod_mlp = self.img_mod(temb).chunk(2, dim=-1)  # [B, 3*dim] each
         txt_mod_attn, txt_mod_mlp = self.txt_mod(temb).chunk(2, dim=-1)  # [B, 3*dim] each
@@ -299,7 +301,8 @@ class QwenImageTransformerBlock(nn.Module):
         img_attn_out, txt_attn_out = self.attn(
             image=img_modulated,
             text=txt_modulated,
-            image_rotary_emb=image_rotary_emb,
+            rotary_emb=rotary_emb,
+            attn_mask=attn_mask,
         )
 
         image = image + img_gate * img_attn_out
@@ -368,13 +371,74 @@ class QwenImageDiT(PreTrainedModel):
         )
         return hidden_states
 
+    def process_entity_masks(
+        self,
+        text: torch.Tensor,
+        text_seq_lens: torch.LongTensor,
+        rotary_emb: Tuple[torch.Tensor, torch.Tensor],
+        video_fhw: List[Tuple[int, int, int]],
+        entity_text: List[torch.Tensor],
+        entity_seq_lens: List[torch.LongTensor],
+        entity_masks: List[torch.Tensor],
+        device: str,
+        dtype: torch.dtype,
+    ):
+        entity_seq_lens = [seq_lens.max().item() for seq_lens in entity_seq_lens]
+        text_seq_lens = entity_seq_lens + [text_seq_lens.max().item()]
+        entity_text = [
+            self.txt_in(self.txt_norm(text[:, :seq_len])) for text, seq_len in zip(entity_text, entity_seq_lens)
+        ]
+        text = torch.cat(entity_text + [text], dim=1)
+
+        entity_txt_freqs = [self.pos_embed(video_fhw, seq_len, device)[1] for seq_len in entity_seq_lens]
+        img_freqs, txt_freqs = rotary_emb
+        txt_freqs = torch.cat(entity_txt_freqs + [txt_freqs], dim=0)
+        rotary_emb = (img_freqs, txt_freqs)
+
+        global_mask = torch.ones_like(entity_masks[0], device=device, dtype=dtype)
+        patched_masks = [self.patchify(mask) for mask in entity_masks + [global_mask]]
+        batch_size, image_seq_len = patched_masks[0].shape[:2]
+        total_seq_len = sum(text_seq_lens) + image_seq_len
+        attention_mask = torch.ones((batch_size, total_seq_len, total_seq_len), device=device, dtype=torch.bool)
+
+        # text-image attention mask
+        img_start, img_end = sum(text_seq_lens), total_seq_len
+        cumsum = [0]
+        for seq_len in text_seq_lens:
+            cumsum.append(cumsum[-1] + seq_len)
+        for i, patched_mask in enumerate(patched_masks):
+            txt_start, txt_end = cumsum[i], cumsum[i + 1]
+            mask = torch.sum(patched_mask, dim=-1) > 0
+            mask = mask.unsqueeze(1).repeat(1, text_seq_lens[i], 1)
+            # text-to-image attention
+            attention_mask[:, txt_start:txt_end, img_start:img_end] = mask
+            # image-to-text attention
+            attention_mask[:, img_start:img_end, txt_start:txt_end] = mask.transpose(1, 2)
+        # entity text tokens should not attend to each other
+        for i in range(len(text_seq_lens)):
+            for j in range(len(text_seq_lens)):
+                if i == j:
+                    continue
+                i_start, i_end = cumsum[i], cumsum[i + 1]
+                j_start, j_end = cumsum[j], cumsum[j + 1]
+                attention_mask[:, i_start:i_end, j_start:j_end] = False
+
+        attn_mask = torch.zeros_like(attention_mask, device=device, dtype=dtype)
+        attn_mask[~attention_mask] = -torch.inf
+        attn_mask = attn_mask.unsqueeze(1)
+        return text, rotary_emb, attn_mask
+
     def forward(
         self,
         image: torch.Tensor,
         edit: torch.Tensor = None,
-        text: torch.Tensor = None,
         timestep: torch.LongTensor = None,
-        txt_seq_lens: torch.LongTensor = None,
+        text: torch.Tensor = None,
+        text_seq_lens: torch.LongTensor = None,
+        context_latents: Optional[torch.Tensor] = None,
+        entity_text: Optional[List[torch.Tensor]] = None,
+        entity_seq_lens: Optional[List[torch.LongTensor]] = None,
+        entity_masks: Optional[List[torch.Tensor]] = None,
     ):
         h, w = image.shape[-2:]
         fp8_linear_enabled = getattr(self, "fp8_linear_enabled", False)
@@ -385,36 +449,62 @@ class QwenImageDiT(PreTrainedModel):
             cfg_parallel(
                 (
                     image,
-                    edit,
-                    text,
+                    *(edit if edit is not None else ()),
                     timestep,
-                    txt_seq_lens,
+                    text,
+                    text_seq_lens,
+                    *(entity_text if entity_text is not None else ()),
+                    *(entity_seq_lens if entity_seq_lens is not None else ()),
+                    *(entity_masks if entity_masks is not None else ()),
+                    context_latents,
                 ),
                 use_cfg=use_cfg,
             ),
         ):
             conditioning = self.time_text_embed(timestep, image.dtype)
             video_fhw = [(1, h // 2, w // 2)]  # frame, height, width
-            max_length = txt_seq_lens.max().item()
+            text_seq_len = text_seq_lens.max().item()
             image = self.patchify(image)
             image_seq_len = image.shape[1]
+            if context_latents is not None:
+                context_latents = context_latents.to(dtype=image.dtype)
+                context_latents = self.patchify(context_latents)
+                image = torch.cat([image, context_latents], dim=1)
+                video_fhw += [(1, h // 2, w // 2)]
             if edit is not None:
-                edit = edit.to(dtype=image.dtype)
-                edit = self.patchify(edit)
-                image = torch.cat([image, edit], dim=1)
-                video_fhw += video_fhw
+                for img in edit:
+                    img = img.to(dtype=image.dtype)
+                    edit_h, edit_w = img.shape[-2:]
+                    img = self.patchify(img)
+                    image = torch.cat([image, img], dim=1)
+                    video_fhw += [(1, edit_h // 2, edit_w // 2)]
 
-            image_rotary_emb = self.pos_embed(video_fhw, max_length, image.device)
+            rotary_emb = self.pos_embed(video_fhw, text_seq_len, image.device)
 
             image = self.img_in(image)
-            text = self.txt_in(self.txt_norm(text[:, :max_length]))
+            text = self.txt_in(self.txt_norm(text[:, :text_seq_len]))
+
+            attn_mask = None
+            if entity_text is not None:
+                text, rotary_emb, attn_mask = self.process_entity_masks(
+                    text,
+                    text_seq_lens,
+                    rotary_emb,
+                    video_fhw,
+                    entity_text,
+                    entity_seq_lens,
+                    entity_masks,
+                    image.device,
+                    image.dtype,
+                )
 
             for block in self.transformer_blocks:
-                text, image = block(image=image, text=text, temb=conditioning, image_rotary_emb=image_rotary_emb)
+                text, image = block(
+                    image=image, text=text, temb=conditioning, rotary_emb=rotary_emb, attn_mask=attn_mask
+                )
             image = self.norm_out(image, conditioning)
             image = self.proj_out(image)
-            if edit is not None:
-                image = image[:, :image_seq_len]
+            image = image[:, :image_seq_len]
 
             image = self.unpatchify(image, h, w)
 

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/models/wan/wan_audio_encoder.py

@@ -267,9 +267,13 @@ def linear_interpolation(features: torch.Tensor, input_fps: int, output_fps: int
     return output_features.transpose(1, 2)  # [1, output_len, 512]
 
 
-def extract_audio_feat(audio_input: torch.Tensor, model: Wav2Vec2Model, dtype=torch.float32, device="cuda:0") -> torch.Tensor:
+def extract_audio_feat(
+    audio_input: torch.Tensor, model: Wav2Vec2Model, dtype=torch.float32, device="cuda:0"
+) -> torch.Tensor:
     video_rate = 30
-    input_values = (audio_input - audio_input.mean(dim=1, keepdim=True)) / torch.sqrt(audio_input.var(dim=1, keepdim=True) + 1e-7)
+    input_values = (audio_input - audio_input.mean(dim=1, keepdim=True)) / torch.sqrt(
+        audio_input.var(dim=1, keepdim=True) + 1e-7
+    )
     feat = torch.cat(model(input_values.to(device)))
     feat = linear_interpolation(feat, input_fps=50, output_fps=video_rate)
     return feat.to(dtype)  # Encoding for the motion

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/pipelines/base.py

@@ -2,10 +2,10 @@ import os
 import torch
 import numpy as np
 from einops import rearrange
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
 from PIL import Image
 
-from diffsynth_engine.configs import BaseConfig, BaseStateDicts
+from diffsynth_engine.configs import BaseConfig, BaseStateDicts, LoraConfig
 from diffsynth_engine.utils.offload import enable_sequential_cpu_offload, offload_model_to_dict, restore_model_from_dict
 from diffsynth_engine.utils.fp8_linear import enable_fp8_autocast
 from diffsynth_engine.utils.gguf import load_gguf_checkpoint
@@ -53,7 +53,7 @@ class BasePipeline:
 
     def update_weights(self, state_dicts: BaseStateDicts) -> None:
         raise NotImplementedError()
-    
+
     @staticmethod
     def update_component(
         component: torch.nn.Module,
@@ -65,10 +65,27 @@ class BasePipeline:
             component.load_state_dict(state_dict, assign=True)
             component.to(device=device, dtype=dtype, non_blocking=True)
 
-    def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
-        for lora_path, lora_scale in lora_list:
-            logger.info(f"loading lora from {lora_path} with scale {lora_scale}")
+    def load_loras(
+        self,
+        lora_list: List[Tuple[str, Union[float, LoraConfig]]],
+        fused: bool = True,
+        save_original_weight: bool = False,
+    ):
+        for lora_path, lora_item in lora_list:
+            if isinstance(lora_item, float):
+                lora_scale = lora_item
+                scheduler_config = None
+            if isinstance(lora_item, LoraConfig):
+                lora_scale = lora_item.scale
+                scheduler_config = lora_item.scheduler_config
+
+            logger.info(f"loading lora from {lora_path} with LoraConfig (scale={lora_scale})")
             state_dict = load_file(lora_path, device=self.device)
+
+            if scheduler_config is not None:
+                self.apply_scheduler_config(scheduler_config)
+                logger.info(f"Applied scheduler args from LoraConfig: {scheduler_config}")
+
             lora_state_dict = self.lora_converter.convert(state_dict)
             for model_name, state_dict in lora_state_dict.items():
                 model = getattr(self, model_name)
@@ -92,6 +109,9 @@ class BasePipeline:
     def load_lora(self, path: str, scale: float, fused: bool = True, save_original_weight: bool = False):
         self.load_loras([(path, scale)], fused, save_original_weight)
 
+    def apply_scheduler_config(self, scheduler_config: Dict):
+        pass
+
     def unload_loras(self):
         raise NotImplementedError()
 

{diffsynth_engine-0.6.1.dev9 → diffsynth_engine-0.6.1.dev11}/diffsynth_engine/pipelines/hunyuan3d_shape.py

@@ -200,5 +200,5 @@ class Hunyuan3DShapePipeline(BasePipeline):
             model_outputs = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
             latents = self.sampler.step(latents, model_outputs, i)
             if progress_callback is not None:
-                progress_callback(i, len(timesteps), "DENOISING")            
+                progress_callback(i, len(timesteps), "DENOISING")
         return self.decode_latents(latents)