Add support for non-divisible batch sizes

src/smc/pipeline.py

@@ -581,6 +581,21 @@ class Pipeline(
         else:
             img_ids = _prepare_latent_image_ids(model_input.shape[0],2*model_input.shape[1],2*model_input.shape[2],model_input.device,model_input.dtype)
         txt_ids = torch.zeros(encoder_hidden_states.shape[1],3).to(device = encoder_hidden_states.device, dtype = encoder_hidden_states.dtype)
+    
+        if prompt_embeds.shape[0] != model_input.shape[0]:
+            # This can happen for the last batch (if batch_p is not divisble by total particles)
+            if guidance_scale > 1.0:
+                batch_p = prompt_embeds.shape[0] // 2
+                last_batch_size = model_input.shape[0] // 2
+                prompt_embeds = torch.cat([prompt_embeds[:last_batch_size], prompt_embeds[batch_p :batch_p + last_batch_size]])
+                encoder_hidden_states = torch.cat([encoder_hidden_states[:last_batch_size], encoder_hidden_states[batch_p :batch_p + last_batch_size]])
+                micro_conds = torch.cat([micro_conds[:last_batch_size], micro_conds[batch_p :batch_p + last_batch_size]])
+            else:
+                last_batch_size = model_input.shape[0]
+                prompt_embeds = prompt_embeds[:last_batch_size]
+                encoder_hidden_states = encoder_hidden_states[:last_batch_size]
+                micro_conds = micro_conds[:last_batch_size]
+        
         model_output = self.transformer(
             hidden_states = model_input,
             micro_conds=micro_conds,