AlekseyCalvin commited on
Commit
0ab9334
1 Parent(s): 160f6b5

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +4 -0
pipeline.py CHANGED
@@ -235,6 +235,10 @@ class FluxWithCFGPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile
235
  scale_lora_layers(self.text_encoder_2, lora_scale)
236
 
237
  prompt = [prompt] if isinstance(prompt, str) else prompt
 
 
 
 
238
 
239
  if prompt_embeds is None:
240
  prompt_2 = prompt_2 or prompt
 
235
  scale_lora_layers(self.text_encoder_2, lora_scale)
236
 
237
  prompt = [prompt] if isinstance(prompt, str) else prompt
238
+ if prompt is not None:
239
+ batch_size = len(prompt)
240
+ else:
241
+ batch_size = prompt_embeds.shape[0]
242
 
243
  if prompt_embeds is None:
244
  prompt_2 = prompt_2 or prompt