Skip to content

Commit 7e0a4c2

Browse files
Support Hunyuan3d (#150)
* add memory predict model by lr * 3d * fix hunyuan bug * ruff format * Update diffsynth_engine/models/hunyuan3d/dino_image_encoder.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update diffsynth_engine/utils/memory/memory_predcit_model.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * bug fix * Update diffsynth_engine/pipelines/hunyuan3d_shape.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 924c6c4 commit 7e0a4c2

File tree

26 files changed

+1803
-16
lines changed

26 files changed

+1803
-16
lines changed

diffsynth_engine/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
SDImagePipeline,
1818
WanVideoPipeline,
1919
QwenImagePipeline,
20+
Hunyuan3DShapePipeline,
2021
)
2122
from .models.flux import FluxControlNet, FluxIPAdapter, FluxRedux
2223
from .models.sd import SDControlNet
@@ -57,6 +58,7 @@
5758
"SDImagePipeline",
5859
"WanVideoPipeline",
5960
"QwenImagePipeline",
61+
"Hunyuan3DShapePipeline",
6062
"FluxInpaintingTool",
6163
"FluxOutpaintingTool",
6264
"FluxIPAdapterRefTool",

diffsynth_engine/algorithm/noise_scheduler/flow_match/recifited_flow.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
import torch
22
import math
33

4-
from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import append_zero, BaseScheduler
4+
from diffsynth_engine.algorithm.noise_scheduler.base_scheduler import BaseScheduler
5+
6+
7+
def append(x, value):
8+
return torch.cat([x, x.new_ones([1]) * value])
59

610

711
class RecifitedFlowScheduler(BaseScheduler):
@@ -37,6 +41,7 @@ def schedule(
3741
mu: float | None = None,
3842
sigma_min: float | None = None,
3943
sigma_max: float | None = None,
44+
append_value: float = 0,
4045
):
4146
sigma_min = self.sigma_min if sigma_min is None else sigma_min
4247
sigma_max = self.sigma_max if sigma_max is None else sigma_max
@@ -46,5 +51,5 @@ def schedule(
4651
else:
4752
sigmas = self._shift_sigma(sigmas, self.shift)
4853
timesteps = sigmas * self.num_train_timesteps
49-
sigmas = append_zero(sigmas)
54+
sigmas = append(sigmas, append_value)
5055
return sigmas, timesteps

diffsynth_engine/algorithm/sampler/flow_match/flow_match_euler.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,14 @@
22

33

44
class FlowMatchEulerSampler:
5-
def initialize(self, init_latents, timesteps, sigmas, mask=None):
6-
self.init_latents = init_latents
7-
self.timesteps = timesteps
5+
def initialize(self, sigmas):
86
self.sigmas = sigmas
9-
self.mask = mask
107

118
def step(self, latents, model_outputs, i):
129
dt = self.sigmas[i + 1] - self.sigmas[i]
1310
latents = latents.to(dtype=torch.float32)
1411
latents = latents + model_outputs * dt
1512
latents = latents.to(dtype=model_outputs.dtype)
16-
if self.mask is not None:
17-
latents = latents * self.mask + self.init_latents * (1 - self.mask)
1813
return latents
1914

2015
def add_noise(self, latents, noise, sigma):

diffsynth_engine/configs/pipeline.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,16 @@ def __post_init__(self):
184184
init_parallel_config(self)
185185

186186

187+
@dataclass
188+
class HunyuanPipelineConfig(BaseConfig):
189+
model_path: str | os.PathLike | List[str | os.PathLike]
190+
model_dtype: torch.dtype = torch.float16
191+
vae_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
192+
vae_dtype: torch.dtype = torch.float16
193+
image_encoder_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
194+
image_encoder_dtype: torch.dtype = torch.float16
195+
196+
187197
@dataclass
188198
class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, BaseConfig):
189199
model_path: str | os.PathLike | List[str | os.PathLike]

diffsynth_engine/models/basic/attention.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ def attention(
111111
q,
112112
k,
113113
v,
114-
attn_impl: Optional[str] = None,
114+
attn_impl: Optional[str] = "auto",
115115
attn_mask: Optional[torch.Tensor] = None,
116116
scale: Optional[float] = None,
117117
**kwargs,
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
from .hunyuan3d_dit import HunYuan3DDiT
2+
from .hunyuan3d_vae import ShapeVAEDecoder
3+
from .dino_image_encoder import ImageEncoder
4+
5+
__all__ = [
6+
"HunYuan3DDiT",
7+
"ShapeVAEDecoder",
8+
"ImageEncoder",
9+
]

0 commit comments

Comments
 (0)