Make flash attn optional. Handle larger batch sizes.

2025-04-14 14:34:46 +00:00
parent 89c0f688db
commit 524bd2edfc
3 changed files with 34 additions and 15 deletions
--- a/extensions_built_in/diffusion_models/hidream/hidream_model.py
+++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py
@@ -129,14 +129,6 @@ class HidreamModel(BaseModel):
            torch_dtype=torch.bfloat16
        )
        
-        # count the params
-        sd = transformer.state_dict()
-        num_params = sum(p.numel() for p in sd.values())
-        print(f"Number of params in transformer: {num_params}")
-        # count params with name expert in them
-        num_expert_params = sum(p.numel() for k, p in sd.items() if 'expert' in k)
-        print(f"Number of params in transformer with expert in them: {num_expert_params}")
-        
        if not self.low_vram:
            transformer.to(self.device_torch, dtype=dtype)