Speed up omnigen TE loading

2026-02-09 23:19:57 +00:00 · 2025-07-05 09:32:00 -06:00
parent 5906a76666
commit bbb57de6ec
1 changed files with 1 additions and 1 deletions
--- 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)