diffusion_utils - Infernet ML
Description: Infernet ML Documentation

Skip to content

# Module: `diffusion_utils`

## `VaeConfig`

Bases: `BaseModel`

Configuration class for VAE (Variational Autoencoder).

Attributes:

Name | Type | Description
--- | --- | ---
`type` | `VaeType` | The type of VAE.
`model` | `str` | The model name.
`torch_dtype` | `Any` | The torch data type.

Source code in `src/infernet_ml/utils/diffusion_utils.py`

    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    class VaeConfig(BaseModel):
        """
        Configuration class for VAE (Variational Autoencoder).

        Attributes:
            type (VaeType): The type of VAE.
            model (str): The model name.
            torch_dtype (Any): The torch data type.
        """

        type: VaeType
        model: str
        torch_dtype: Any  # Change torch.dtype to Any

        class Config:
            arbitrary_types_allowed = True  # Add this line to allow arbitrary types

## `get_torch_dtype(dtype)`

Returns the corresponding torch.dtype based on the input string dtype.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`dtype` | `str` | The input string representing the desired dtype. | required

Returns:

Type | Description
--- | ---
`dtype` | torch.dtype: The corresponding torch.dtype.

Examples:

    >>> get_torch_dtype("float16")
    torch.float16

    >>> get_torch_dtype("float32")
    torch.float32

    >>> get_torch_dtype("int")
    torch.float32

Source code in `src/infernet_ml/utils/diffusion_utils.py`

    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    def get_torch_dtype(dtype: str) -> torch.dtype:
        """
        Returns the corresponding torch.dtype based on the input string dtype.

        Args:
            dtype (str): The input string representing the desired dtype.

        Returns:
            torch.dtype: The corresponding torch.dtype.

        Raises:
            None

        Examples:
            >>> get_torch_dtype("float16")
            torch.float16

            >>> get_torch_dtype("float32")
            torch.float32

            >>> get_torch_dtype("int")
            torch.float32
        """
        match dtype:
            case "float16" | "half":
                return torch.float16
            case "float32" | "single":
                return torch.float32
            case "float64" | "double":
                return torch.float64
            case _:
                return torch.float32

## `get_vae_kl(config)`

Retrieves a Variational Autoencoder (VAE) model based on the given configuration.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`config` | `VaeConfig` | The configuration object specifying the type of VAE model to retrieve. | required

Returns:

Name | Type | Description
--- | --- | ---
`Any` | `Any` | The VAE model.

Examples:

    >>> config = VaeConfig(type=VaeType.ema, torch_dtype=torch.float32)
    >>> vae = get_vae_kl(config)

Source code in `src/infernet_ml/utils/diffusion_utils.py`

    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    def get_vae_kl(config: VaeConfig) -> Any:
        """
        Retrieves a Variational Autoencoder (VAE) model based on the given configuration.

        Args:
            config (VaeConfig): The configuration object specifying the type of VAE
              model to retrieve.

        Returns:
            Any: The VAE model.

        Raises:
            None

        Examples:
            >>> config = VaeConfig(type=VaeType.ema, torch_dtype=torch.float32)
            >>> vae = get_vae_kl(config)
        """
        vae_ft_ema = "stabilityai/sd-vae-ft-ema"
        vae_ft_mse = "stabilityai/sd-vae-ft-mse"
        match config.type:
            case VaeType.ema:
                vae = AutoencoderKL.from_pretrained(  # type: ignore
                    vae_ft_ema, torch_dtype=config.torch_dtype
                )
            case VaeType.mse:
                vae = AutoencoderKL.from_pretrained(  # type: ignore
                    vae_ft_mse, torch_dtype=config.torch_dtype
                )
            case _:
                vae = AutoencoderKL.from_pretrained(
                    config.model, subfolder="vae", torch_dtype=config.torch_dtype
                )
        return vae.to("cuda")