refactor(flux2): make block slab config backend neutral

Author: Watebear

lightx2v/models/networks/flux2/model.py

@@ -14,7 +14,6 @@
     Flux2TransformerWeights,
     preserve_weight_module_cpu_tensors,
     release_weight_module_device_tensors,
-    validate_flux2_block_slab_config,
 )
 from lightx2v.utils.custom_compiler import compiled_method
 from lightx2v_platform.base import global_var
@@ -28,10 +27,10 @@ class _Flux2TransformerModelBase(BaseTransformerModel):
 
     def __init__(self, config, model_path, device):
         super().__init__(model_path, config, device)
-        self.use_block_slab_offload = validate_flux2_block_slab_config(
-            self.config,
-            global_var.AI_DEVICE,
-        )
+        # Block-slab packing is platform-agnostic. Enable it only on backends that support
+        # event offload, with block-level CPU offload, BF16 Default weights, and no LoRA,
+        # lazy loading, or tensor parallelism.
+        self.use_block_slab_offload = self.config.get("offload_use_block_slab", False)
         self._offload_weights_loaded = False
         self._offload_weights_preparing = False
         self.in_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 64))

lightx2v/models/networks/flux2/weights/transformer_weights.py

@@ -7,32 +7,6 @@
 from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER
 
 
-def validate_flux2_block_slab_config(config, ai_device):
-    """Validate the deliberately narrow first rollout of block slabs."""
-    if not config.get("offload_use_block_slab", False):
-        return False
-
-    requirements = []
-    if ai_device != "npu":
-        requirements.append("AI_DEVICE='npu'")
-    if not config.get("cpu_offload", False) or config.get("offload_granularity", "block") != "block":
-        requirements.append("block-level cpu_offload")
-    if not config.get("use_event_offload", False):
-        requirements.append("use_event_offload=true")
-    if config.get("dit_quantized", False) or config.get("dit_quant_scheme", "Default") != "Default":
-        requirements.append("unquantized dit_quant_scheme='Default'")
-    if config.get("lora_configs") or config.get("lora_dynamic_apply", False):
-        requirements.append("LoRA disabled")
-    if config.get("lazy_load", False):
-        requirements.append("lazy_load=false")
-    if config.get("tensor_parallel", False):
-        requirements.append("tensor_parallel=false")
-
-    if requirements:
-        raise ValueError("offload_use_block_slab currently requires " + ", ".join(requirements))
-    return True
-
-
 def _resolve_resident_block_indices(value, num_blocks, policy, config_key):
     """Resolve a resident-block count into deterministic block indices.