From 35588726de6de1fb870740d11d1f428c79ab800b Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 7 Jun 2026 13:31:03 -0600 Subject: [PATCH] Fixed issue where a buffer was stuck on cpu when offloading ideogram4 --- extensions_built_in/diffusion_models/ideogram4/ideogram4.py | 2 +- .../diffusion_models/ideogram4/src/transformer.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index c3c53a85..810cd99d 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -289,7 +289,7 @@ class Ideogram4Model(BaseModel): transformer, self.device_torch, offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[transformer.input_proj, transformer.llm_cond_proj], + ignore_modules=[transformer.rotary_emb.inv_freq, transformer.input_proj, transformer.llm_cond_proj], ) elif self.model_config.low_vram: self.print_and_status_update("Moving transformer to CPU") diff --git a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py index dd594293..cb146fc6 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py +++ b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py @@ -93,6 +93,10 @@ class Ideogram4MRoPE(nn.Module): # position_ids: (B, L, 3) of int. assert position_ids.ndim == 3 and position_ids.shape[-1] == 3 batch_size, seq_len, _ = position_ids.shape + + if self.inv_freq.device == torch.device("cpu"): + # sometimes it gets stuck on CPU + self.inv_freq = self.inv_freq.to(position_ids.device) # (3, B, inv_freq_size, L) pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)