def __init__()

in docker_images/latent-to-image/app/pipelines/latent_to_image.py [0:0]

• 24 lines of code
• 8 McCabe index (conditional complexity)

    def __init__(self, model_id: str):
        self.model_id = None
        self.current_tokens_loaded = 0
        self.use_auth_token = os.getenv("HF_API_TOKEN")
        # This should allow us to make the image work with private models when no token is provided, if the said model
        # is already in local cache
        self.offline_preferred = validation.str_to_bool(os.getenv("OFFLINE_PREFERRED"))
        model_data = self._hub_model_info(model_id)

        kwargs = {}
        env_dtype = os.getenv("TORCH_DTYPE", "float32")
        if env_dtype:
            kwargs["torch_dtype"] = getattr(torch, env_dtype)
        elif torch.cuda.is_available():
            kwargs["torch_dtype"] = torch.float16

        has_model_index = any(
            file.rfilename == "model_index.json" for file in model_data.siblings
        )

        if has_model_index:
            kwargs["subfolder"] = "vae"

        self.vae = AutoencoderKL.from_pretrained(model_id, **kwargs).eval()
        self.dtype = kwargs["torch_dtype"]
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)

        if not idle.UNLOAD_IDLE:
            self._model_to_gpu()