2323from PIL import Image
2424from transformers import AutoModel , AutoModelForCausalLM , AutoTokenizer
2525
26- from ...utils .torch_utils import randn_tensor
2726from ...models import AutoencoderKLFlux2
2827from ...models .transformers import ErnieImageTransformer2DModel
2928from ...pipelines .pipeline_utils import DiffusionPipeline
3029from ...schedulers import FlowMatchEulerDiscreteScheduler
30+ from ...utils .torch_utils import randn_tensor
3131from .pipeline_output import ErnieImagePipelineOutput
3232
3333
@@ -206,8 +206,8 @@ def __call__(
206206 num_images_per_prompt : int = 1 ,
207207 generator : Optional [torch .Generator ] = None ,
208208 latents : Optional [torch .Tensor ] = None ,
209- prompt_embeds : list [torch .FloatTensor ] | None = None ,
210- negative_prompt_embeds : list [torch .FloatTensor ] | None = None ,
209+ prompt_embeds : list [torch .FloatTensor ] | None = None ,
210+ negative_prompt_embeds : list [torch .FloatTensor ] | None = None ,
211211 output_type : str = "pil" ,
212212 return_dict : bool = True ,
213213 callback_on_step_end : Optional [Callable [[int , int , dict ], None ]] = None ,
@@ -267,10 +267,7 @@ def __call__(
267267 # [Phase 1] PE: enhance prompts
268268 revised_prompts : Optional [List [str ]] = None
269269 if prompt is not None and use_pe and self .pe is not None and self .pe_tokenizer is not None :
270- prompt = [
271- self ._enhance_prompt_with_pe (p , device , width = width , height = height )
272- for p in prompt
273- ]
270+ prompt = [self ._enhance_prompt_with_pe (p , device , width = width , height = height ) for p in prompt ]
274271 revised_prompts = list (prompt )
275272
276273 if prompt is not None :
@@ -298,9 +295,7 @@ def __call__(
298295 if negative_prompt_embeds is not None :
299296 uncond_text_hiddens = negative_prompt_embeds
300297 else :
301- uncond_text_hiddens = self .encode_prompt (
302- negative_prompt , device , num_images_per_prompt
303- )
298+ uncond_text_hiddens = self .encode_prompt (negative_prompt , device , num_images_per_prompt )
304299
305300 # Latent dimensions
306301 latent_h = height // self .vae_scale_factor
@@ -310,10 +305,10 @@ def __call__(
310305 # Initialize latents
311306 if latents is None :
312307 latents = randn_tensor (
313- (total_batch_size , latent_channels , latent_h , latent_w ),
314- generator = generator ,
315- device = device ,
316- dtype = dtype
308+ (total_batch_size , latent_channels , latent_h , latent_w ),
309+ generator = generator ,
310+ device = device ,
311+ dtype = dtype ,
317312 )
318313
319314 # Setup scheduler
0 commit comments