More progress on brainstorming code for asset management for models

comfy_extras/nodes_assets_test.py

@@ -0,0 +1,56 @@
+from comfy_api.latest import io, ComfyExtension
+import comfy.asset_management
+import comfy.sd
+import folder_paths
+import logging
+import os
+
+
+class AssetTestNode(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="AssetTestNode",
+            is_experimental=True,
+            inputs=[
+                io.Combo.Input("ckpt_name", folder_paths.get_filename_list("checkpoints")),
+            ],
+            outputs=[
+                io.Model.Output(),
+                io.Clip.Output(),
+                io.Vae.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, ckpt_name: str):
+        hash = None
+        # lets get the full path just so we can retrieve the hash from db, if exists
+        try:
+            full_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+            if full_path is None:
+                raise Exception(f"Model {ckpt_name} not found")
+            from app.model_processor import model_processor
+            hash = model_processor.retrieve_hash(full_path)
+        except Exception as e:
+            logging.error(f"Could not get model by hash with error: {e}")
+        subdir, name = os.path.split(ckpt_name)
+        asset_info = comfy.asset_management.AssetInfo(hash=hash, name=name, tags=["models", "checkpoints"], metadata={"subdir": subdir})
+        asset = comfy.asset_management.resolve(asset_info)
+        # /\ the stuff above should happen in execution code instead of inside the node
+        # \/ the stuff below should happen in the node - confirm is a model asset, do stuff to it (already loaded? or should be called to 'load'?)
+        if asset is None:
+            raise Exception(f"Model {asset_info.name} not found")
+        assert isinstance(asset, comfy.asset_management.ModelReturnedAsset)
+        out = comfy.sd.load_state_dict_guess_config(asset.state_dict, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"), metadata=asset.metadata)
+        return io.NodeOutput(out[0], out[1], out[2])
+
+
+class AssetTestExtension(ComfyExtension):
+    @classmethod
+    async def get_node_list(cls):
+        return [AssetTestNode]
+
+
+def comfy_entrypoint():
+    return AssetTestExtension()