PixArt initial rewrite

2024-12-10 22:19:02 +01:00
parent 7dc1403cd5
commit 3f844cac42
28 changed files with 583 additions and 3066 deletions
@@ -0,0 +1,34 @@
+import folder_paths
+import comfy.utils
+
+from .PixArt.loader import load_pixart_state_dict
+
+loaders = {
+    "PixArt": load_pixart_state_dict,
+}
+
+class EXMUnetLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "unet_name": (folder_paths.get_filename_list("unet"),),
+                "model_type": (list(loaders.keys()),)
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load_unet"
+    CATEGORY = "ExtraModels"
+    TITLE = "Load Diffusion Model (ExtraModels)"
+
+    def load_unet(self, unet_name, model_type):
+        model_options = {}
+        unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
+        loader_fn = loaders[model_type]
+        sd = comfy.utils.load_torch_file(unet_path)
+        return (loader_fn(sd),)
+
+NODE_CLASS_MAPPINGS = {
+    "EXMUnetLoader": EXMUnetLoader,
+}