Fixed issue where a buffer was stuck on cpu when offloading ideogram4

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")

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)