Function encode_prompt

train_code/train_wan_motion.py:322–345

(
    tokenizer: AutoTokenizer,
    text_encoder: UMT5EncoderModel,
    prompt: Union[str, List[str]],
    num_videos_per_prompt: int = 1,
    max_sequence_length: int = 226,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    text_input_ids=None,
)

def encode_prompt(
    tokenizer: AutoTokenizer,
    text_encoder: UMT5EncoderModel,
    prompt: Union[str, List[str]],
    num_videos_per_prompt: int = 1,
    max_sequence_length: int = 226,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    text_input_ids=None,
):
    
    # prompt = [prompt] if isinstance(prompt, str) else prompt
    prompt_embeds = _get_t5_prompt_embeds(
                                            tokenizer,
                                            text_encoder,
                                            prompt = prompt,
                                            num_videos_per_prompt = num_videos_per_prompt,
                                            max_sequence_length = max_sequence_length,
                                            device = device,
                                            dtype = dtype,
                                            text_input_ids = text_input_ids,
                                        )

    return prompt_embeds