diff --git a/README.md b/README.md index 0fd64c6..17bd63d 100644 --- a/README.md +++ b/README.md @@ -53,11 +53,11 @@ python main.py - `Model Download Batch`: 追加・削除ボタンのある一覧UI。内部ではcanonical JSONを保持します。 - `Model Download Batch (JSON)`: 同じコアを使うJSON直接入力版です。 -- `Downloaded Checkpoint Loader`: Illustriousなどの一体型checkpointをロードします。 -- `Downloaded Diffusion Model Loader`: AnimaやKrea 2のdiffusion modelをロードします。 -- `Downloaded Text Encoder Loader`: Animaでは`auto`、Krea 2では`krea2`を選びます。 -- `Downloaded VAE Loader`: Qwen Image VAEなどをロードします。 -- `Downloaded LoRA Loader`: 対応LoRAをロードします。model-only LoRAでは`strength_clip`を0にできます。 +- `Load Checkpoint (Downloaded)`: Illustriousなどの一体型checkpointをロードします。 +- `Load Diffusion Model (Downloaded)`: AnimaやKrea 2のdiffusion modelをロードします。 +- `Load CLIP (Downloaded)`: ComfyUI標準の`Load CLIP`と同じ`type`の選択肢を使用します。`krea2`、`ideogram4`、`flux2`などを選択できます。 +- `Load VAE (Downloaded)`: Qwen Image VAEなどをロードします。 +- `Load LoRA (Downloaded)`: 対応LoRAをロードします。model-only LoRAでは`strength_clip`を0にできます。 2種類のdownload nodeは、任意型の`passthrough`を入力と出力に持ちます。既存の接続途中に挿入できるほか、未接続でもoutput nodeとして実行できます。取得結果を使う場合は`download_result`を対応するDownloaded Loaderへ接続し、manifestの`id`を指定します。 diff --git a/model_batch_downloader/loaders.py b/model_batch_downloader/loaders.py index eb34041..b9d67e9 100644 --- a/model_batch_downloader/loaders.py +++ b/model_batch_downloader/loaders.py @@ -62,17 +62,18 @@ class ModelBatchDownloaderDiffusionModelLoader: return UNETLoader().load_unet(record.relative_path, weight_dtype) -class ModelBatchDownloaderTextEncoderLoader: +class ModelBatchDownloaderCLIPLoader: @classmethod def INPUT_TYPES(cls): + standard = CLIPLoader.INPUT_TYPES() return { "required": { "download_result": ("DOWNLOAD_RESULT",), "id": ("STRING", {"default": ""}), - "type": (["auto", "krea2"],), + "type": standard["required"]["type"], }, "optional": { - "device": (["default", "cpu"], {"advanced": True}), + "device": standard["optional"]["device"], }, } @@ -82,8 +83,7 @@ class ModelBatchDownloaderTextEncoderLoader: def load(self, download_result, id, type, device="default"): record = _record(download_result, id, "text_encoders") - clip_type = "stable_diffusion" if type == "auto" else type - return CLIPLoader().load_clip(record.relative_path, clip_type, device) + return CLIPLoader().load_clip(record.relative_path, type, device) class ModelBatchDownloaderVAELoader: @@ -151,15 +151,15 @@ class ModelBatchDownloaderLoRALoader: LOADER_CLASS_MAPPINGS = { "ModelBatchDownloaderCheckpointLoader": ModelBatchDownloaderCheckpointLoader, "ModelBatchDownloaderDiffusionModelLoader": ModelBatchDownloaderDiffusionModelLoader, - "ModelBatchDownloaderTextEncoderLoader": ModelBatchDownloaderTextEncoderLoader, + "ModelBatchDownloaderCLIPLoader": ModelBatchDownloaderCLIPLoader, "ModelBatchDownloaderVAELoader": ModelBatchDownloaderVAELoader, "ModelBatchDownloaderLoRALoader": ModelBatchDownloaderLoRALoader, } LOADER_DISPLAY_NAME_MAPPINGS = { - "ModelBatchDownloaderCheckpointLoader": "Downloaded Checkpoint Loader", - "ModelBatchDownloaderDiffusionModelLoader": "Downloaded Diffusion Model Loader", - "ModelBatchDownloaderTextEncoderLoader": "Downloaded Text Encoder Loader", - "ModelBatchDownloaderVAELoader": "Downloaded VAE Loader", - "ModelBatchDownloaderLoRALoader": "Downloaded LoRA Loader", + "ModelBatchDownloaderCheckpointLoader": "Load Checkpoint (Downloaded)", + "ModelBatchDownloaderDiffusionModelLoader": "Load Diffusion Model (Downloaded)", + "ModelBatchDownloaderCLIPLoader": "Load CLIP (Downloaded)", + "ModelBatchDownloaderVAELoader": "Load VAE (Downloaded)", + "ModelBatchDownloaderLoRALoader": "Load LoRA (Downloaded)", } diff --git a/tests/test_loaders.py b/tests/test_loaders.py index f2ed9bc..6224cdd 100644 --- a/tests/test_loaders.py +++ b/tests/test_loaders.py @@ -2,6 +2,17 @@ from model_batch_downloader.aria2_runner import DownloadRecord, DownloadResult import model_batch_downloader.loaders as subject +subject.CLIPLoader.INPUT_TYPES = classmethod( + lambda cls: { + "required": { + "clip_name": (["model.safetensors"],), + "type": (["stable_diffusion", "krea2", "ideogram4", "flux2"],), + }, + "optional": {"device": (["default", "cpu"], {"advanced": True})}, + } +) + + def result(model_type, relative_path="folder/model.safetensors", item_id="asset"): return DownloadResult( { @@ -44,20 +55,26 @@ def test_diffusion_loader_passes_weight_dtype(monkeypatch): assert seen == [("folder/model.safetensors", "fp8_e4m3fn")] -def test_text_encoder_maps_auto_and_krea2(monkeypatch): +def test_clip_loader_reuses_standard_type_and_device_inputs(): + standard = subject.CLIPLoader.INPUT_TYPES() + custom = subject.ModelBatchDownloaderCLIPLoader.INPUT_TYPES() + assert custom["required"]["type"] == standard["required"]["type"] + assert custom["optional"]["device"] == standard["optional"]["device"] + assert "auto" not in custom["required"]["type"][0] + + +def test_clip_loader_delegates_type_and_device(monkeypatch): seen = [] monkeypatch.setattr( subject.CLIPLoader, "load_clip", lambda self, name, type, device: seen.append((name, type, device)) or ("clip",), ) - loader = subject.ModelBatchDownloaderTextEncoderLoader() - loader.load(result("text_encoders"), "asset", "auto", "default") - loader.load(result("text_encoders"), "asset", "krea2", "cpu") - assert seen == [ - ("folder/model.safetensors", "stable_diffusion", "default"), - ("folder/model.safetensors", "krea2", "cpu"), - ] + output = subject.ModelBatchDownloaderCLIPLoader().load( + result("text_encoders"), "asset", "ideogram4", "cpu" + ) + assert output == ("clip",) + assert seen == [("folder/model.safetensors", "ideogram4", "cpu")] def test_wrong_category_fails_before_delegation(): diff --git a/tests/test_registration.py b/tests/test_registration.py index fe692f2..3bb2756 100644 --- a/tests/test_registration.py +++ b/tests/test_registration.py @@ -17,11 +17,22 @@ def test_plugin_registers_seven_namespaced_nodes(): assert sorted(module.NODE_CLASS_MAPPINGS) == [ "ModelBatchDownloader", + "ModelBatchDownloaderCLIPLoader", "ModelBatchDownloaderCheckpointLoader", "ModelBatchDownloaderDiffusionModelLoader", "ModelBatchDownloaderJSON", "ModelBatchDownloaderLoRALoader", - "ModelBatchDownloaderTextEncoderLoader", "ModelBatchDownloaderVAELoader", ] + assert "ModelBatchDownloaderCLIPLoader" in module.NODE_CLASS_MAPPINGS + assert "ModelBatchDownloaderTextEncoderLoader" not in module.NODE_CLASS_MAPPINGS + assert module.NODE_DISPLAY_NAME_MAPPINGS == { + "ModelBatchDownloader": "Model Download Batch", + "ModelBatchDownloaderJSON": "Model Download Batch (JSON)", + "ModelBatchDownloaderCheckpointLoader": "Load Checkpoint (Downloaded)", + "ModelBatchDownloaderDiffusionModelLoader": "Load Diffusion Model (Downloaded)", + "ModelBatchDownloaderCLIPLoader": "Load CLIP (Downloaded)", + "ModelBatchDownloaderVAELoader": "Load VAE (Downloaded)", + "ModelBatchDownloaderLoRALoader": "Load LoRA (Downloaded)", + } assert module.WEB_DIRECTORY == "./web"