improve type inference

model/loader.py

@@ -51,7 +51,10 @@ class Dict2Model:
         setattr(utils, 'load_torch_file', load_torch_file_hook)
         
         try:
-            return sd.load_checkpoint(config_path, None, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
+            model, clip, vae = sd.load_checkpoint(config_path, None, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
+            assert clip is not None
+            assert vae is not None
+            return (model, clip, vae)
         finally:
             setattr(sd, 'load_torch_file', load_torch_file_org)
 

model/merge.py

@@ -11,7 +11,7 @@ def merge(
     half: str,
     ignore_keys_only_in_B: bool = False,
 ):
-    result = dict()
+    result: Dict[str,torch.Tensor] = dict()
     for key in tqdm.tqdm(model_A.keys()):
         if key not in model_B:
             print(f'  key {key} is found in model_A but not model_B')