From 1a039913134e0f67147f62b30d421954f4e9b8a0 Mon Sep 17 00:00:00 2001 From: Haoming Date: Thu, 28 May 2026 11:03:19 +0800 Subject: [PATCH] typing --- backend/diffusion_engine/base.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/backend/diffusion_engine/base.py b/backend/diffusion_engine/base.py index 540f3525..17fca707 100644 --- a/backend/diffusion_engine/base.py +++ b/backend/diffusion_engine/base.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from backend.patcher.clip import CLIP from backend.patcher.unet import UnetPatcher from backend.patcher.vae import VAE + from modules_forge.packages.huggingface_guess.model_list import BASE from backend import memory_management, utils @@ -25,11 +26,11 @@ class ForgeObjects: class ForgeDiffusionEngine: matched_guesses = [] - def __init__(self, estimated_config, huggingface_components): + def __init__(self, estimated_config: "BASE", huggingface_components: dict[str, "torch.nn.Module"]): huggingface_components["vae"].latent_format = estimated_config.latent_format - self.model_config = estimated_config - self.is_inpaint = estimated_config.inpaint_model() + self.model_config: "BASE" = estimated_config + self.is_inpaint: bool = estimated_config.inpaint_model() self.forge_objects: "ForgeObjects" = None self.forge_objects_original: "ForgeObjects" = None