diffsynth-engine 0.4.2.dev6__tar.gz → 0.4.3__tar.gz

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/.gitignore

@@ -8,4 +8,5 @@ dist/
 *.egg-info/
 .DS_Store/
 .pytest_cache/
-.ruff_cache/
+.ruff_cache/
+CLAUDE.md

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.2.dev6
+Version: 0.4.3
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
@@ -25,13 +25,14 @@ Requires-Dist: imageio[ffmpeg]
 Requires-Dist: yunchang; sys_platform == "linux"
 Requires-Dist: onnxruntime
 Requires-Dist: opencv-python
+Requires-Dist: scikit-image
+Requires-Dist: trimesh
 Provides-Extra: dev
 Requires-Dist: diffusers==0.31.0; extra == "dev"
-Requires-Dist: transformers==4.45.2; extra == "dev"
+Requires-Dist: transformers==4.52.4; extra == "dev"
 Requires-Dist: accelerate; extra == "dev"
 Requires-Dist: build; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
-Requires-Dist: scikit-image; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 Dynamic: license-file

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/diffsynth_engine/__init__.py

@@ -4,6 +4,10 @@ from .configs import (
     FluxPipelineConfig,
     WanPipelineConfig,
     QwenImagePipelineConfig,
+    SDStateDicts,
+    SDXLStateDicts,
+    FluxStateDicts,
+    QwenImageStateDicts,
     ControlNetParams,
     ControlType,
 )
@@ -13,11 +17,18 @@ from .pipelines import (
     SDImagePipeline,
     WanVideoPipeline,
     QwenImagePipeline,
+    Hunyuan3DShapePipeline,
 )
 from .models.flux import FluxControlNet, FluxIPAdapter, FluxRedux
 from .models.sd import SDControlNet
 from .models.sdxl import SDXLControlNetUnion
-from .utils.download import fetch_model, fetch_modelscope_model, fetch_civitai_model
+from .utils.download import (
+    fetch_model,
+    fetch_modelscope_model,
+    fetch_civitai_model,
+    register_fetch_modelscope_model,
+    reset_fetch_modelscope_model,
+)
 from .utils.video import load_video, save_video
 from .tools import (
     FluxInpaintingTool,
@@ -32,6 +43,10 @@ __all__ = [
     "SDXLPipelineConfig",
     "FluxPipelineConfig",
     "WanPipelineConfig",
+    "SDStateDicts",
+    "SDXLStateDicts",
+    "FluxStateDicts",
+    "QwenImageStateDicts",
     "FluxImagePipeline",
     "QwenImagePipelineConfig",
     "FluxControlNet",
@@ -43,6 +58,7 @@ __all__ = [
     "SDImagePipeline",
     "WanVideoPipeline",
     "QwenImagePipeline",
+    "Hunyuan3DShapePipeline",
     "FluxInpaintingTool",
     "FluxOutpaintingTool",
     "FluxIPAdapterRefTool",
@@ -52,6 +68,8 @@ __all__ = [
     "ControlType",
     "fetch_model",
     "fetch_modelscope_model",
+    "register_fetch_modelscope_model",
+    "reset_fetch_modelscope_model",
     "fetch_civitai_model",
     "load_video",
     "save_video",

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/diffsynth_engine/algorithm/noise_scheduler/flow_match/recifited_flow.py

@@ -1,7 +1,11 @@
 import torch
 import math
 
-from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero, BaseScheduler
+from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import BaseScheduler
+
+
+def append(x, value):
+    return torch.cat([x, x.new_ones([1]) * value])
 
 
 class RecifitedFlowScheduler(BaseScheduler):
@@ -37,6 +41,7 @@ class RecifitedFlowScheduler(BaseScheduler):
         mu: float | None = None,
         sigma_min: float | None = None,
         sigma_max: float | None = None,
+        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
@@ -46,5 +51,5 @@ class RecifitedFlowScheduler(BaseScheduler):
         else:
             sigmas = self._shift_sigma(sigmas, self.shift)
         timesteps = sigmas * self.num_train_timesteps
-        sigmas = append_zero(sigmas)
+        sigmas = append(sigmas, append_value)
         return sigmas, timesteps

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/diffsynth_engine/algorithm/sampler/flow_match/flow_match_euler.py

@@ -2,19 +2,14 @@ import torch
 
 
 class FlowMatchEulerSampler:
-    def initialize(self, init_latents, timesteps, sigmas, mask=None):
-        self.init_latents = init_latents
-        self.timesteps = timesteps
+    def initialize(self, sigmas):
         self.sigmas = sigmas
-        self.mask = mask
 
     def step(self, latents, model_outputs, i):
         dt = self.sigmas[i + 1] - self.sigmas[i]
         latents = latents.to(dtype=torch.float32)
         latents = latents + model_outputs * dt
         latents = latents.to(dtype=model_outputs.dtype)
-        if self.mask is not None:
-            latents = latents * self.mask + self.init_latents * (1 - self.mask)
         return latents
 
     def add_noise(self, latents, noise, sigma):

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/diffsynth_engine/conf/models/qwen_image/qwen2_5_vl_config.json

@@ -21,5 +21,6 @@
     "vision_start_token_id": 151652,
     "vision_end_token_id": 151653,
     "image_token_id": 151655,
-    "video_token_id": 151656
+    "video_token_id": 151656,
+    "attn_impl": "sdpa"
 }

diffsynth_engine-0.4.3/diffsynth_engine/conf/tokenizers/qwen_image/qwen2_vl_image_processor.json

@@ -0,0 +1,29 @@
+{
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+        0.48145466,
+        0.4578275,
+        0.40821073
+    ],
+    "image_processor_type": "Qwen2VLImageProcessor",
+    "image_std": [
+        0.26862954,
+        0.26130258,
+        0.27577711
+    ],
+    "max_pixels": 12845056,
+    "merge_size": 2,
+    "min_pixels": 3136,
+    "patch_size": 14,
+    "processor_class": "Qwen2_5_VLProcessor",
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+        "longest_edge": 12845056,
+        "shortest_edge": 3136
+    },
+    "temporal_patch_size": 2
+}

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/diffsynth_engine/configs/pipeline.py

@@ -16,6 +16,7 @@ class BaseConfig:
     vae_tile_stride: int | Tuple[int, int] = 256
     device: str = "cuda"
     offload_mode: Optional[str] = None
+    offload_to_disk: bool = False
 
 
 @dataclass
@@ -62,11 +63,13 @@ class SDPipelineConfig(BaseConfig):
         model_path: str | os.PathLike | List[str | os.PathLike],
         device: str = "cuda",
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "SDPipelineConfig":
         return cls(
             model_path=model_path,
             device=device,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
 
@@ -87,11 +90,13 @@ class SDXLPipelineConfig(BaseConfig):
         model_path: str | os.PathLike | List[str | os.PathLike],
         device: str = "cuda",
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "SDXLPipelineConfig":
         return cls(
             model_path=model_path,
             device=device,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
 
@@ -116,6 +121,7 @@ class FluxPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Ba
         device: str = "cuda",
         parallelism: int = 1,
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "FluxPipelineConfig":
         return cls(
             model_path=model_path,
@@ -123,6 +129,7 @@ class FluxPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Ba
             parallelism=parallelism,
             use_fsdp=True,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
     def __post_init__(self):
@@ -160,6 +167,7 @@ class WanPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Bas
         device: str = "cuda",
         parallelism: int = 1,
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "WanPipelineConfig":
         return cls(
             model_path=model_path,
@@ -169,12 +177,23 @@ class WanPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Bas
             use_cfg_parallel=True,
             use_fsdp=True,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
     def __post_init__(self):
         init_parallel_config(self)
 
 
+@dataclass
+class HunyuanPipelineConfig(BaseConfig):
+    model_path: str | os.PathLike | List[str | os.PathLike]
+    model_dtype: torch.dtype = torch.float16
+    vae_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
+    vae_dtype: torch.dtype = torch.float16
+    image_encoder_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
+    image_encoder_dtype: torch.dtype = torch.float16
+
+
 @dataclass
 class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, BaseConfig):
     model_path: str | os.PathLike | List[str | os.PathLike]
@@ -187,6 +206,11 @@ class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfi
     # override OptimizationConfig
     fbcache_relative_l1_threshold = 0.009
 
+    # override BaseConfig
+    vae_tiled: bool = True
+    vae_tile_size: Tuple[int, int] = (34, 34)
+    vae_tile_stride: Tuple[int, int] = (18, 16)
+
     @classmethod
     def basic_config(
         cls,
@@ -196,6 +220,7 @@ class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfi
         device: str = "cuda",
         parallelism: int = 1,
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "QwenImagePipelineConfig":
         return cls(
             model_path=model_path,
@@ -206,6 +231,7 @@ class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfi
             use_cfg_parallel=True,
             use_fsdp=True,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
     def __post_init__(self):

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3}/diffsynth_engine/models/basic/attention.py

@@ -1,9 +1,9 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from einops import rearrange, repeat
 from typing import Optional
 
-import torch.nn.functional as F
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.flag import (
     FLASH_ATTN_3_AVAILABLE,
@@ -42,11 +42,11 @@ if XFORMERS_AVAILABLE:
 
 if SDPA_AVAILABLE:
 
-    def sdpa_attn(q, k, v, attn_mask=None, scale=None):
+    def sdpa_attn(q, k, v, attn_mask=None, is_causal=False, scale=None):
         q = q.transpose(1, 2)
         k = k.transpose(1, 2)
         v = v.transpose(1, 2)
-        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal, scale=scale)
         return out.transpose(1, 2)
 
 
@@ -111,7 +111,7 @@ def attention(
     q,
     k,
     v,
-    attn_impl: Optional[str] = None,
+    attn_impl: Optional[str] = "auto",
     attn_mask: Optional[torch.Tensor] = None,
     scale: Optional[float] = None,
     **kwargs,

diffsynth_engine-0.4.3/diffsynth_engine/models/hunyuan3d/__init__.py

@@ -0,0 +1,9 @@
+from .hunyuan3d_dit import HunYuan3DDiT
+from .hunyuan3d_vae import ShapeVAEDecoder
+from .dino_image_encoder import ImageEncoder
+
+__all__ = [
+    "HunYuan3DDiT",
+    "ShapeVAEDecoder",
+    "ImageEncoder",
+]