Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-13 10:08:43 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix typing

Browse Source
This commit is contained in:
awsr
2026-04-23 16:49:50 -07:00
parent 09f5fc6e03
commit e9d74cea72
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/file_loader.py
+4 -4
View File
@@ -3,7 +3,7 @@ import concurrent.futures
import torch
def map_keys(key: str, key_mapping: dict) -> str:
def map_keys(key: str, key_mapping: dict | None) -> str:
new_key = key
if key_mapping:
for pattern, replacement in key_mapping.items():
@@ -13,7 +13,7 @@ def map_keys(key: str, key_mapping: dict) -> str:
return new_key
def load_safetensors(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> dict:
def load_safetensors(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> None:
from safetensors.torch import safe_open
if state_dict is None:
state_dict = {}
@@ -23,7 +23,7 @@ def load_safetensors(files: list[str], state_dict: dict | None = None, key_mappi
state_dict[map_keys(key, key_mapping)] = f.get_tensor(key)
def load_threaded(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> dict:
def load_threaded(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> None:
future_items = {}
if state_dict is None:
state_dict = {}
@@ -34,7 +34,7 @@ def load_threaded(files: list[str], state_dict: dict | None = None, key_mapping:
future.result()
def load_streamer(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> dict:
def load_streamer(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> None:
# requires pip install runai_model_streamer
from runai_model_streamer import SafetensorsStreamer
if state_dict is None: