@@ -281,6 +281,86 @@ def test_batch_with_num_images_per_prompt(self):
281281 # Should return 4 images (2 prompts × 2 images per prompt)
282282 self .assertEqual (len (images ), 4 )
283283
284+ def test_prompt_with_prior_token_ids (self ):
285+ """Test that prompt and prior_token_ids can be provided together.
286+
287+ When both are given, the AR generation step is skipped (prior_token_ids is used
288+ directly) and prompt is used to generate prompt_embeds via the glyph encoder.
289+ """
290+ device = "cpu"
291+ components = self .get_dummy_components ()
292+ pipe = self .pipeline_class (** components )
293+ pipe .to (device )
294+ pipe .set_progress_bar_config (disable = None )
295+
296+ height , width = 32 , 32
297+
298+ # Step 1: Run with prompt only to get prior_token_ids from AR model
299+ generator = torch .Generator (device = device ).manual_seed (0 )
300+ prior_token_ids , _ , _ = pipe .generate_prior_tokens (
301+ prompt = "A photo of a cat" ,
302+ height = height ,
303+ width = width ,
304+ device = torch .device (device ),
305+ generator = torch .Generator (device = device ).manual_seed (0 ),
306+ )
307+
308+ # Step 2: Run with both prompt and prior_token_ids — should not raise
309+ generator = torch .Generator (device = device ).manual_seed (0 )
310+ inputs_both = {
311+ "prompt" : "A photo of a cat" ,
312+ "prior_token_ids" : prior_token_ids ,
313+ "generator" : generator ,
314+ "num_inference_steps" : 2 ,
315+ "guidance_scale" : 1.5 ,
316+ "height" : height ,
317+ "width" : width ,
318+ "max_sequence_length" : 16 ,
319+ "output_type" : "pt" ,
320+ }
321+ images = pipe (** inputs_both ).images
322+ self .assertEqual (len (images ), 1 )
323+ self .assertEqual (images [0 ].shape , (3 , 32 , 32 ))
324+
325+ def test_check_inputs_rejects_invalid_combinations (self ):
326+ """Test that check_inputs correctly rejects invalid input combinations."""
327+ device = "cpu"
328+ components = self .get_dummy_components ()
329+ pipe = self .pipeline_class (** components )
330+ pipe .to (device )
331+
332+ height , width = 32 , 32
333+
334+ # Neither prompt nor prior_token_ids → error
335+ with self .assertRaises (ValueError ):
336+ pipe .check_inputs (
337+ prompt = None ,
338+ height = height ,
339+ width = width ,
340+ callback_on_step_end_tensor_inputs = None ,
341+ prompt_embeds = torch .randn (1 , 16 , 32 ),
342+ )
343+
344+ # prior_token_ids alone without prompt or prompt_embeds → error
345+ with self .assertRaises (ValueError ):
346+ pipe .check_inputs (
347+ prompt = None ,
348+ height = height ,
349+ width = width ,
350+ callback_on_step_end_tensor_inputs = None ,
351+ prior_token_ids = torch .randint (0 , 100 , (1 , 64 )),
352+ )
353+
354+ # prompt + prompt_embeds together → error
355+ with self .assertRaises (ValueError ):
356+ pipe .check_inputs (
357+ prompt = "A cat" ,
358+ height = height ,
359+ width = width ,
360+ callback_on_step_end_tensor_inputs = None ,
361+ prompt_embeds = torch .randn (1 , 16 , 32 ),
362+ )
363+
284364 @unittest .skip ("Needs to be revisited." )
285365 def test_encode_prompt_works_in_isolation (self ):
286366 pass
0 commit comments