From bbb57de6ec14eae821c08e2fa636158d8e279d18 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 5 Jul 2025 09:32:00 -0600 Subject: [PATCH] Speed up omnigen TE loading --- extensions_built_in/diffusion_models/omnigen2/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index ee37a1f0..ddff811e 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -84,7 +84,7 @@ class OmniGen2Model(BaseModel): subfolder="mllm", torch_dtype=torch.bfloat16 ) - + mllm.to(self.device_torch, dtype=dtype) if self.model_config.quantize_te: self.print_and_status_update("Quantizing Qwen2.5 VL model") quantization_type = get_qtype(self.model_config.qtype_te)