Function sample_image_inference

library/lumina_train_util.py:250–478

Generates sample images Args: accelerator (Accelerator): Accelerator object args (argparse.Namespace): Arguments object nextdit (lumina_models.NextDiT): NextDiT model gemma2_model (list[Gemma2Model]): Gemma2 model vae (AutoEncoder): VAE model

(
    accelerator: Accelerator,
    args: argparse.Namespace,
    nextdit: lumina_models.NextDiT,
    gemma2_model: list[Gemma2Model],
    vae: AutoEncoder,
    save_dir: str,
    prompt_dicts: list[Dict[str, str]],
    epoch: int,
    global_step: int,
    sample_prompts_gemma2_outputs: dict[str, Tuple[Tensor, Tensor, Tensor]],
    prompt_replacement: Optional[Tuple[str, str]] = None,
    controlnet=None,
)

@torch.no_grad()
def sample_image_inference(
    accelerator: Accelerator,
    args: argparse.Namespace,
    nextdit: lumina_models.NextDiT,
    gemma2_model: list[Gemma2Model],
    vae: AutoEncoder,
    save_dir: str,
    prompt_dicts: list[Dict[str, str]],
    epoch: int,
    global_step: int,
    sample_prompts_gemma2_outputs: dict[str, Tuple[Tensor, Tensor, Tensor]],
    prompt_replacement: Optional[Tuple[str, str]] = None,
    controlnet=None,
):
    """
    Generates sample images

    Args:
        accelerator (Accelerator): Accelerator object
        args (argparse.Namespace): Arguments object
        nextdit (lumina_models.NextDiT): NextDiT model
        gemma2_model (list[Gemma2Model]): Gemma2 model
        vae (AutoEncoder): VAE model
        save_dir (str): Directory to save images
        prompt_dict (Dict[str, str]): Prompt dictionary
        epoch (int): Epoch number
        steps (int): Number of steps to run
        sample_prompts_gemma2_outputs (List[Tuple[Tensor, Tensor, Tensor]]): List of tuples containing Gemma 2 outputs
        prompt_replacement (Optional[Tuple[str, str]], optional): Replacement for positive and negative prompt. Defaults to None.

    Returns:
        None
    """

    # encode prompts
    tokenize_strategy = strategy_base.TokenizeStrategy.get_strategy()
    encoding_strategy = strategy_base.TextEncodingStrategy.get_strategy()

    assert isinstance(tokenize_strategy, strategy_lumina.LuminaTokenizeStrategy)
    assert isinstance(encoding_strategy, strategy_lumina.LuminaTextEncodingStrategy)

    text_conds = []

    # assuming seed, width, height, sample steps, guidance are the same
    width = int(prompt_dicts[0].get("width", 1024))
    height = int(prompt_dicts[0].get("height", 1024))
    height = max(64, height - height % 8)  # round to divisible by 8
    width = max(64, width - width % 8)  # round to divisible by 8

    guidance_scale = float(prompt_dicts[0].get("scale", 3.5))
    cfg_trunc_ratio = float(prompt_dicts[0].get("cfg_trunc_ratio", 0.25))
    renorm_cfg = float(prompt_dicts[0].get("renorm_cfg", 1.0))
    sample_steps = int(prompt_dicts[0].get("sample_steps", 36))
    seed = prompt_dicts[0].get("seed", None)
    seed = int(seed) if seed is not None else None
    assert seed is None or seed > 0, f"Invalid seed {seed}"
    generator = torch.Generator(device=accelerator.device)
    if seed is not None: