Function main_kwargs

diffusers/examples/research_projects/vae/vae_roundtrip.py:153–190

(
    *,
    device: torch.device,
    input_image_path: str,
    pretrained_model_name_or_path: str,
    revision: Optional[str],
    variant: Optional[str],
    subfolder: Optional[str],
    use_tiny_nn: bool,
)

def main_kwargs(
    *,
    device: torch.device,
    input_image_path: str,
    pretrained_model_name_or_path: str,
    revision: Optional[str],
    variant: Optional[str],
    subfolder: Optional[str],
    use_tiny_nn: bool,
) -> None:
    vae = load_vae_model(
        device=device,
        model_name_or_path=pretrained_model_name_or_path,
        revision=revision,
        variant=variant,
        subfolder=subfolder,
        use_tiny_nn=use_tiny_nn,
    )
    original_pil = Image.open(input_image_path).convert("RGB")
    original_image = pil_to_nhwc(
        device=device,
        image=original_pil,
    )
    print(f"Original image shape: {original_image.shape}")
    reconstructed_image: Optional[torch.Tensor] = None

    with torch.no_grad():
        latent_image = to_latent(rgb_nchw=original_image, vae=vae)
        print(f"Latent shape: {latent_image.shape}")
        reconstructed_image = from_latent(latent_nchw=latent_image, vae=vae)
        reconstructed_pil = nhwc_to_pil(nhwc=reconstructed_image)
    combined_image = concatenate_images(
        left=original_pil,
        right=reconstructed_pil,
        vertical=False,
    )
    combined_image.show("Original | Reconstruction")
    print(f"Reconstructed image shape: {reconstructed_image.shape}")


def parse_args() -> argparse.Namespace: