Function _enhance

gemma_encoder.py:402–448

Common enhancement logic for both T2V and I2V modes.

(
    processor: Gemma3Processor,
    model: Gemma3ForConditionalGeneration,
    messages: list,
    image: Optional[Image.Image] = None,
    max_new_tokens: int = 512,
    seed: int = 42,
)

def _enhance(
    processor: Gemma3Processor,
    model: Gemma3ForConditionalGeneration,
    messages: list,
    image: Optional[Image.Image] = None,
    max_new_tokens: int = 512,
    seed: int = 42,
) -> str:
    """Common enhancement logic for both T2V and I2V modes."""
    if processor is None:
        raise ValueError("Processor not loaded - enhancement not available")

    text = processor.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    model_inputs = processor(
        text=text,
        images=image,
        return_tensors="pt",
    ).to(model.device)

    pad_token_id = (
        processor.tokenizer.pad_token_id
        if processor.tokenizer.pad_token_id is not None
        else 0
    )
    model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id)

    with (
        torch.inference_mode(),
        torch.random.fork_rng(devices=[model.device]),
        torch.autocast(device_type=model.device.type, dtype=model.dtype),
    ):
        torch.manual_seed(seed)
        outputs = model.generate(
            **model_inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
        )
        generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
        enhanced_prompt = processor.tokenizer.decode(
            generated_ids, skip_special_tokens=True
        )

    return enhanced_prompt


def enhance_t2v(