Merge pull request #4858 from awsr/typing-merge-1

Additional typing updates

modules/errorlimiter.py

@@ -2,17 +2,15 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from threading import Lock
-from typing import TYPE_CHECKING, ClassVar
+from typing import ClassVar
 
-if TYPE_CHECKING:
-    from collections.abc import Iterable
 
 _instance_id = 0
 _lock = Lock()
 
 
 def _make_unique(name: str):
-    global _instance_id # pylint: disable=global-statement
+    global _instance_id  # pylint: disable=global-statement
     with _lock:  # Guard against race conditions
         new_name = f"{name}__{_instance_id}"
         _instance_id += 1

modules/extra_networks.py

@@ -1,18 +1,24 @@
+from __future__ import annotations
+
 import re
 import inspect
 from collections import defaultdict
+from typing import TYPE_CHECKING, cast
 from modules import errors, shared
 from modules.logger import log
 
+if TYPE_CHECKING:
+    from modules.processing_class import StableDiffusionProcessing
 
-extra_network_registry = {}
+
+extra_network_registry: dict[str, ExtraNetwork] = {}
 
 
 def initialize():
     extra_network_registry.clear()
 
 
-def register_extra_network(extra_network):
+def register_extra_network(extra_network: ExtraNetwork):
     extra_network_registry[extra_network.name] = extra_network
 
 
@@ -26,10 +32,10 @@ def register_default_extra_networks():
 
 
 class ExtraNetworkParams:
-    def __init__(self, items=None):
-        self.items = items or []
-        self.positional = []
-        self.named = {}
+    def __init__(self, items: list[str] | None = None):
+        self.items: list[str] = items or []
+        self.positional: list[str] = []
+        self.named: dict[str, str] = {}
         for item in self.items:
             parts = item.split('=', 2) if isinstance(item, str) else [item]
             if len(parts) == 2:
@@ -48,10 +54,10 @@ class ExtraNetworkParams:
 
 
 class ExtraNetwork:
-    def __init__(self, name):
+    def __init__(self, name: str):
         self.name = name
 
-    def activate(self, p, params_list):
+    def activate(self, p: StableDiffusionProcessing, params_list: list[ExtraNetworkParams], *args, **kwargs):
         """
         Called by processing on every run. Whatever the extra network is meant to do should be activated here. Passes arguments related to this extra network in params_list. User passes arguments by specifying this in his prompt:
         <name:arg1:arg2:arg3>
@@ -68,14 +74,14 @@ class ExtraNetwork:
         """
         raise NotImplementedError
 
-    def deactivate(self, p, force=False):
+    def deactivate(self, p: StableDiffusionProcessing, force=False):
         """
         Called at the end of processing for housekeeping. No need to do anything here.
         """
         raise NotImplementedError
 
 
-def activate(p, extra_network_data=None, step=0, include=None, exclude=None):
+def activate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, step=0, include: list | None = None, exclude: list | None = None):
     """call activate for extra networks in extra_network_data in specified order, then call activate for all remaining registered networks with an empty argument list"""
     if exclude is None:
         exclude = []
@@ -115,12 +121,12 @@ def activate(p, extra_network_data=None, step=0, include=None, exclude=None):
     p.network_data = extra_network_data
 
 
-def deactivate(p, extra_network_data=None, force=None):
+def deactivate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, force: bool | None = None):
     """call deactivate for extra networks in extra_network_data in specified order, then call deactivate for all remaining registered networks"""
     if p.disable_extra_networks:
         return
     if force is None:
-        force = shared.opts.lora_force_reload
+        force = cast("bool", shared.opts.lora_force_reload)
     extra_network_data = extra_network_data or p.network_data
 
     for extra_network_name in extra_network_data:
@@ -150,7 +156,7 @@ def parse_prompt(prompt: str | None) -> tuple[str, defaultdict[str, list[ExtraNe
     if prompt is None:
         return "", res
     if isinstance(prompt, list):
-        return parse_prompts(prompt)
+        return parse_prompts(prompt)  # type: ignore --- Fallback for incorrect function calls
 
     def found(m: re.Match[str]):
         name, args = m.group(1, 2)
@@ -161,9 +167,9 @@ def parse_prompt(prompt: str | None) -> tuple[str, defaultdict[str, list[ExtraNe
     return updated_prompt, res
 
 
-def parse_prompts(prompts: list[str], extra_data=None):
+def parse_prompts(prompts: list[str], extra_data: defaultdict[str, list[ExtraNetworkParams]] | None = None):
     updated_prompt_list: list[str] = []
-    extra_data: defaultdict[str, list[ExtraNetworkParams]] = extra_data or defaultdict(list)
+    extra_data = extra_data or defaultdict(list)
     for prompt in prompts:
         updated_prompt, parsed_extra_data = parse_prompt(prompt)
         if not extra_data:

modules/processing_class.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 import os
 import sys
 import inspect
 import hashlib
-from typing import Any
+from collections import defaultdict
+from typing import Any, TYPE_CHECKING
 from dataclasses import dataclass, field
 import numpy as np
 from PIL import Image, ImageOps
@@ -11,6 +14,9 @@ from modules.logger import log
 from modules.paths import resolve_output_path
 from modules.image.util import flatten
 
+if TYPE_CHECKING:
+    from modules.extra_networks import ExtraNetworkParams
+
 
 debug = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
 
@@ -323,7 +329,9 @@ class StableDiffusionProcessing:
         self.negative_prompt_attention_masks = []
         self.disable_extra_networks = False
         self.iteration = 0
-        self.network_data = network_data or {}
+        self.network_data: defaultdict[str, list[ExtraNetworkParams]] = defaultdict(list)
+        if network_data is not None:
+            self.network_data |= network_data
 
         # initializers
         self.prompt = prompt

modules/processing_prompt.py

@@ -1,8 +1,14 @@
+from __future__ import annotations
+
 import os
+from typing import TYPE_CHECKING
 import torch
 from modules import shared, errors, timer, prompt_parser_diffusers
 from modules.logger import log
 
+if TYPE_CHECKING:
+    from modules.processing_class import StableDiffusionProcessing
+
 
 debug_enabled = os.environ.get('SD_PROMPT_DEBUG', None) is not None
 debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
@@ -46,11 +52,11 @@ def fix_prompt_batch(p, prompts, negative_prompts, prompts_2, negative_prompts_2
     return prompts, negative_prompts, prompts_2, negative_prompts_2
 
 
-def fix_prompt_model(cls, prompts, negative_prompts, prompts_2, negative_prompts_2):
+def fix_prompt_model(cls: str, prompts: list, negative_prompts: list, prompts_2: list, negative_prompts_2: list):
     if 'OmniGen' in cls:
         prompts = [p.replace('|image|', '<img><|image_1|></img>') for p in prompts]
     if 'PixArtSigmaPipeline' in cls: # pixart-sigma pipeline throws list-of-list for negative prompt
-        negative_prompts = negative_prompts[0]
+        negative_prompts = negative_prompts[0]  # type: ignore --- Handle list-of-list
     return prompts, negative_prompts, prompts_2, negative_prompts_2
 
 
@@ -70,7 +76,7 @@ def set_fallback_prompt(args: dict, possible: list[str], prompts, negative_promp
     return args
 
 
-def set_prompt(p,
+def set_prompt(p: StableDiffusionProcessing,
                args: dict,
                possible: list[str],
                cls: str,
@@ -81,7 +87,7 @@ def set_prompt(p,
                negative_prompts: list[str],
                prompts_2: list[str],
                negative_prompts_2: list[str],
-              ) -> dict:
+              ):
     prompt_attention = prompt_attention or getattr(p, 'prompt_attention', None) or shared.opts.prompt_attention
     if (prompt_attention != 'fixed') and ('Onnx' not in cls) and ('prompt' not in p.task_args) and (
         ('StableDiffusion' in cls) or

modules/sd_checkpoint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 import base64
 import os
@@ -11,8 +13,8 @@ from modules.logger import log
 from modules.json_helpers import writefile
 
 
-checkpoints_list = {}
-checkpoint_aliases = {}
+checkpoints_list: dict[str, CheckpointInfo] = {}
+checkpoint_aliases: dict[str, CheckpointInfo] = {}
 checkpoints_loaded = collections.OrderedDict()
 model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
@@ -24,7 +26,7 @@ warn_once = False
 
 
 class CheckpointInfo:
-    def __init__(self, filename, name=None, sha=None, subfolder=None, model_type: str = 'checkpoint', folder: str|None = None):
+    def __init__(self, filename: str, name: str | None = None, sha: str | None = None, subfolder: str | None = None, model_type: str = 'checkpoint', folder: str | None = None):
         self.name = name
         self.hash = sha
         self.filename = filename
@@ -33,7 +35,7 @@ class CheckpointInfo:
         relname = filename
         app_path = os.path.abspath(paths.script_path)
 
-        def rel(fn, path):
+        def rel(fn: str, path: str):
             try:
                 return os.path.relpath(fn, path)
             except Exception:
@@ -208,7 +210,7 @@ def remove_hash(s):
     return re.sub(r'\s*\[.*?\]', '', s)
 
 
-def get_closest_checkpoint_match(s: str) -> CheckpointInfo:
+def get_closest_checkpoint_match(s: str) -> CheckpointInfo | None:
     # direct hf url
     if s.startswith('https://huggingface.co/'):
         model_name = s.replace('https://huggingface.co/', '')
@@ -357,7 +359,7 @@ def extract_thumbnail(filename, data):
         log.error(f"Error extracting thumbnail: {filename} {e}")
 
 
-def read_metadata_from_safetensors(filename):
+def read_metadata_from_safetensors(filename: str):
     global sd_metadata # pylint: disable=global-statement
     if sd_metadata is None:
         sd_metadata = shared.readfile(sd_metadata_file, lock=True, as_type="dict") if os.path.isfile(sd_metadata_file) else {}
@@ -414,7 +416,7 @@ def read_metadata_from_safetensors(filename):
     return res
 
 
-def scrub_dict(dict_obj, keys):
+def scrub_dict(dict_obj, keys: list[str]):
     for key in list(dict_obj.keys()):
         if not isinstance(dict_obj, dict):
             continue

modules/sd_vae.py

@@ -1,15 +1,22 @@
+from __future__ import annotations
+
 import os
 import glob
+from typing import TYPE_CHECKING, cast
 import torch
 from modules import shared, errors, paths, devices, sd_models, sd_detect
 from modules.logger import log
 
+if TYPE_CHECKING:
+    from diffusers import DiffusionPipeline
+    from modules.sd_checkpoint import CheckpointInfo
+
 
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}
-vae_dict = {}
-base_vae = None
-loaded_vae_file = None
-checkpoint_info = None
+vae_dict: dict[str, str] = {}
+base_vae = None  # Unused
+loaded_vae_file: str | None = None
+checkpoint_info: CheckpointInfo | None = None
 vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE'))
 debug = os.environ.get('SD_VAE_DEBUG', None) is not None
 unspecified = object()
@@ -19,7 +26,7 @@ vae_scale_override = {
 }
 
 
-def get_vae_scale_factor(model=None):
+def get_vae_scale_factor(model: DiffusionPipeline | None = None):
     if not shared.sd_loaded:
         vae_scale_factor = 8
         return vae_scale_factor
@@ -41,20 +48,20 @@ def get_vae_scale_factor(model=None):
     else:
         # log.warning(f'VAE: cls={model.__class__.__name__ if model else "None"} scale=unknown')
         vae_scale_factor = 8
-    if hasattr(model, 'patch_size'):
+    if model is not None and hasattr(model, 'patch_size'):
         patch_size = model.patch_size
     if debug:
         log.trace(f'VAE: cls={model.__class__.__name__ if model else "None"} scale={vae_scale_factor} patch={patch_size}')
     return vae_scale_factor * patch_size
 
 
-def load_vae_dict(filename):
+def load_vae_dict(filename: str):
     vae_ckpt = sd_models.read_state_dict(filename, what='vae')
     vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
     return vae_dict_1
 
 
-def get_filename(filepath):
+def get_filename(filepath: str):
     if filepath.endswith(".json"):
         return os.path.basename(os.path.dirname(filepath))
     else:
@@ -92,7 +99,7 @@ def refresh_vae_list():
     return vae_dict
 
 
-def find_vae_near_checkpoint(checkpoint_file):
+def find_vae_near_checkpoint(checkpoint_file: str):
     checkpoint_path = os.path.splitext(checkpoint_file)[0]
     for vae_location in [f"{checkpoint_path}.vae.pt", f"{checkpoint_path}.vae.ckpt", f"{checkpoint_path}.vae.safetensors"]:
         if os.path.isfile(vae_location):
@@ -100,11 +107,11 @@ def find_vae_near_checkpoint(checkpoint_file):
     return None
 
 
-def resolve_vae(checkpoint_file):
+def resolve_vae(checkpoint_file: str):
     if shared.opts.sd_vae == 'TAESD':
         return None, None
     if shared.cmd_opts.vae is not None: # 1st
-        return shared.cmd_opts.vae, 'forced'
+        return cast("str", shared.cmd_opts.vae), 'forced'
     if shared.opts.sd_vae == "Default": # 2nd
         return None, None
     vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file)
@@ -125,7 +132,7 @@ def resolve_vae(checkpoint_file):
     return None, None
 
 
-def apply_vae_config(model_file, vae_file, sd_model):
+def apply_vae_config(model_file: str, vae_file: str, sd_model: DiffusionPipeline):
     def get_vae_config():
         config_file = os.path.join(paths.sd_configs_path, os.path.splitext(os.path.basename(model_file))[0] + '_vae.json')
         if config_file is not None and os.path.exists(config_file):
@@ -145,7 +152,7 @@ def apply_vae_config(model_file, vae_file, sd_model):
                 sd_model.vae.config[k] = v
 
 
-def load_vae(model_file, vae_file=None, vae_source="unknown-source"):
+def load_vae(model_file: str, vae_file: str | None = None, vae_source: str | None = "unknown-source"):
     if vae_file is None:
         return None
     if not os.path.exists(vae_file):
@@ -174,7 +181,7 @@ def load_vae(model_file, vae_file=None, vae_source="unknown-source"):
         import diffusers
         vae_class = None
         vae_loader = None
-        if shared.sd_loaded and getattr(shared.sd_model, 'vae', None) is not None:
+        if shared.sd_model is not None and getattr(shared.sd_model, 'vae', None) is not None:
             vae_class = shared.sd_model.vae.__class__
             vae_loader = vae_class.from_single_file if os.path.isfile(vae_file) else vae_class.from_pretrained
         elif os.path.isfile(vae_file):
@@ -213,7 +220,7 @@ def load_vae(model_file, vae_file=None, vae_source="unknown-source"):
     return None
 
 
-def reload_vae_weights(sd_model=None, vae_file=unspecified):
+def reload_vae_weights(sd_model: DiffusionPipeline | None = None, vae_file = unspecified):
     if not sd_model:
         sd_model = shared.sd_model
     if sd_model is None:
@@ -222,25 +229,27 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
     checkpoint_info = sd_model.sd_checkpoint_info
     checkpoint_file = checkpoint_info.filename
     if vae_file == unspecified:
-        vae_file, vae_source = resolve_vae(checkpoint_file)
+        vae_file_path, vae_source = resolve_vae(checkpoint_file)
     else:
+        vae_file_path = cast("str | None", vae_file)
         vae_source = "function-argument"
-    if vae_file is None or vae_file == 'None':
+
+    if vae_file_path is None or vae_file_path == 'None':
         if hasattr(sd_model, 'original_vae'):
             sd_models.set_diffuser_options(sd_model, vae=sd_model.original_vae, op='vae')
             log.info("VAE restored")
             return None
-    if loaded_vae_file == vae_file:
+    if loaded_vae_file == vae_file_path:
         return None
 
-    if hasattr(sd_model, "vae") and getattr(sd_model, "sd_checkpoint_info", None) is not None:
-        vae = load_vae(sd_model.sd_checkpoint_info.filename, vae_file, vae_source)
+    if vae_file_path is not None and hasattr(sd_model, "vae") and getattr(sd_model, "sd_checkpoint_info", None) is not None:
+        vae = load_vae(sd_model.sd_checkpoint_info.filename, vae_file_path, vae_source)
         if vae is not None:
             if not hasattr(sd_model, 'original_vae'):
                 sd_model.original_vae = sd_model.vae
                 sd_models.move_model(sd_model.original_vae, devices.cpu)
             sd_models.set_diffuser_options(sd_model, vae=vae, op='vae')
-            apply_vae_config(sd_model.sd_checkpoint_info.filename, vae_file, sd_model)
+            apply_vae_config(sd_model.sd_checkpoint_info.filename, vae_file_path, sd_model)
 
     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
         sd_models.move_model(sd_model, devices.device)

modules/theme.py

@@ -1,10 +1,18 @@
-import os
 import json
+import os
+from typing import TYPE_CHECKING, TypedDict
+
 import gradio as gr
-import modules.shared
+
 import modules.extensions
-from modules.logger import log
+import modules.shared
 from modules.json_helpers import writefile
+from modules.logger import log
+
+if TYPE_CHECKING:
+    class FontParams(TypedDict):
+        font: list[str]
+        font_mono: list[str]
 
 
 gradio_theme = gr.themes.Base()
@@ -89,7 +97,7 @@ def list_themes():
 def reload_gradio_theme():
     global gradio_theme # pylint: disable=global-statement
     theme_name = modules.shared.opts.gradio_theme
-    default_font_params = {
+    default_font_params: FontParams = {
         'font':['Helvetica', 'ui-sans-serif', 'system-ui', 'sans-serif'],
         'font_mono':['IBM Plex Mono', 'ui-monospace', 'Consolas', 'monospace']
     }