github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fixed issue with training hidream when batch size is larger than 1

Browse Source
This commit is contained in:
Jaret Burkett
2025-04-21 17:26:29 +00:00
parent 12e3095d8a
commit add83df5cc
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
extensions_built_in/diffusion_models/hidream/hidream_model.py
View File

@@ -329,6 +329,7 @@ class HidreamModel(BaseModel):
text_embeddings: PromptEmbeds, text_embeddings: PromptEmbeds,
**kwargs **kwargs
): ):
batch_size = latent_model_input.shape[0]
with torch.no_grad(): with torch.no_grad():
if latent_model_input.shape[-2] != latent_model_input.shape[-1]: if latent_model_input.shape[-2] != latent_model_input.shape[-1]:
B, C, H, W = latent_model_input.shape B, C, H, W = latent_model_input.shape
@@ -343,7 +344,9 @@ class HidreamModel(BaseModel):
img_ids_pad[:pH*pW, :] = img_ids img_ids_pad[:pH*pW, :] = img_ids
img_sizes = img_sizes.unsqueeze(0).to(latent_model_input.device) img_sizes = img_sizes.unsqueeze(0).to(latent_model_input.device)
img_sizes = torch.cat([img_sizes] * batch_size, dim=0)
img_ids = img_ids_pad.unsqueeze(0).to(latent_model_input.device) img_ids = img_ids_pad.unsqueeze(0).to(latent_model_input.device)
img_ids = torch.cat([img_ids] * batch_size, dim=0)
else: else:
img_sizes = img_ids = None img_sizes = img_ids = None