ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 11:11:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Work on mean flow. Minor bug fixes. Omnigen improvements

Browse Source
This commit is contained in:
Jaret Burkett
2025-06-26 13:46:20 -06:00
parent 84c6edca7e
commit 8d9c47316a
4 changed files with 128 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
extensions_built_in/diffusion_models/omnigen2/__init__.py
View File

@@ -239,7 +239,10 @@ class OmniGen2Model(BaseModel):
**kwargs
):
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep = timestep.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
try:
timestep = timestep.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
except Exception as e:
pass
# optional_kwargs = {}
# if 'ref_image_hidden_states' in set(inspect.signature(self.model.forward).parameters.keys()):