From 3ea5ef986502c52a239aac302ebb3204bc8092e5 Mon Sep 17 00:00:00 2001 From: 2kpr <96332338+2kpr@users.noreply.github.com> Date: Wed, 17 Jul 2024 12:49:36 -0500 Subject: [PATCH] Add custom dirs --- README.md | 2 + nodes.py | 64 ++++++++++++-- ultrapixel.py | 54 +++++++++--- workflow_controlnet.json | 178 ++++++++++++++++++++------------------- workflow_default.json | 108 ++++++++++++------------ 5 files changed, 248 insertions(+), 158 deletions(-) diff --git a/README.md b/README.md index ba1e228..f0e741d 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,8 @@ When the workflow first runs the first node will download all the necessary file To enable ControlNet usage you merely have to use the load image node in ComfyUI and tie that to the controlnet_image input on the UltraPixel Process node, you can also attach a preview/save image node to the edge_preview output of the UltraPixel Process node to see the controlnet edge preview. Easiest to just load the included workflow_controlnet.json file in ComfyUI. +As mentioned above the default directory for the UltraPixel and StableCascade downloaded model files is ComfyUI/models/ultrapixel, if you want to alter this you can now change ultrapixel_directory or stablecascade_directory in the UltraPixel Load node from 'default' to the full path/directory you desire. + Example Output for prompt: "A close-up portrait of a young woman with flawless skin, vibrant red lipstick, and wavy brown hair, wearing a vintage floral dress and standing in front of a blooming garden."
diff --git a/nodes.py b/nodes.py index b4981c5..e5b553f 100644 --- a/nodes.py +++ b/nodes.py @@ -16,7 +16,7 @@ def get_models(base): return models -def download_models(): +def download_models(ultrapixel_directory, stablecascade_directory): models = [ ["ultrapixel_t2i.safetensors", "2kpr/UltraPixel"], ["stage_a.safetensors", "stabilityai/stable-cascade"], @@ -27,7 +27,28 @@ def download_models(): ["controlnet/canny.safetensors", "stabilityai/stable-cascade"], ] for model in models: - model_path = os.path.join(folder_paths.models_dir, "ultrapixel") + if model[1].startswith("stabilityai"): + if stablecascade_directory == "default" or not os.path.exists( + stablecascade_directory + ): + model_path = os.path.join(folder_paths.models_dir, "ultrapixel") + if not os.path.exists(stablecascade_directory): + print( + f"{stablecascade_directory} does not exist, defaulting to {model_path}" + ) + else: + model_path = stablecascade_directory + else: + if ultrapixel_directory == "default" or not os.path.exists( + ultrapixel_directory + ): + model_path = os.path.join(folder_paths.models_dir, "ultrapixel") + if not os.path.exists(ultrapixel_directory): + print( + f"{ultrapixel_directory} does not exist, defaulting to {model_path}" + ) + else: + model_path = ultrapixel_directory if os.path.exists(os.path.join(model_path, model[0])): # print("Path already exists:", os.path.join(model_path, model[0])) continue @@ -51,6 +72,22 @@ class UltraPixelLoad: "effnet": (get_models("effnet_encoder.safetensors"),), "previewer": (get_models("previewer.safetensors"),), "controlnet": (get_models("controlnet/canny.safetensors"),), + "ultrapixel_directory": ( + "STRING", + { + "multiline": False, + "dynamicPrompts": True, + "default": "default", + }, + ), + "stablecascade_directory": ( + "STRING", + { + "multiline": False, + "dynamicPrompts": True, + "default": "default", + }, + ), }, } @@ -60,11 +97,28 @@ class UltraPixelLoad: CATEGORY = "UltraPixel" def ultrapixel( - self, ultrapixel, stage_a, stage_b, stage_c, effnet, previewer, controlnet + self, + ultrapixel, + stage_a, + stage_b, + stage_c, + effnet, + previewer, + controlnet, + ultrapixel_directory, + stablecascade_directory, ): - download_models() + download_models(ultrapixel_directory, stablecascade_directory) model = up.UltraPixel( - ultrapixel, stage_a, stage_b, stage_c, effnet, previewer, controlnet + ultrapixel, + stage_a, + stage_b, + stage_c, + effnet, + previewer, + controlnet, + ultrapixel_directory, + stablecascade_directory, ) return (model,) diff --git a/ultrapixel.py b/ultrapixel.py index 6921910..355eff8 100644 --- a/ultrapixel.py +++ b/ultrapixel.py @@ -19,16 +19,48 @@ from .train import WurstCore_t2i as WurstCoreC class UltraPixel: def __init__( - self, pretrained, stage_a, stage_b, stage_c, effnet, previewer, controlnet + self, + pretrained, + stage_a, + stage_b, + stage_c, + effnet, + previewer, + controlnet, + ultrapixel_directory, + stablecascade_directory, ): - self.ultrapixel_path = os.path.join(folder_paths.models_dir, "ultrapixel") + if ultrapixel_directory == "default" or not os.path.exists( + ultrapixel_directory + ): + self.ultrapixel_path = os.path.join(folder_paths.models_dir, "ultrapixel") + if not os.path.exists(ultrapixel_directory): + print( + f"{ultrapixel_directory} does not exist, defaulting to {self.ultrapixel_path}" + ) + else: + self.ultrapixel_path = ultrapixel_directory + if stablecascade_directory == "default" or not os.path.exists( + stablecascade_directory + ): + self.stablecascade_path = os.path.join( + folder_paths.models_dir, "ultrapixel" + ) + if not os.path.exists(stablecascade_directory): + print( + f"{stablecascade_directory} does not exist, defaulting to {self.stablecascade_path}" + ) + else: + self.stablecascade_path = stablecascade_directory self.pretrained = os.path.join(self.ultrapixel_path, pretrained) - self.stage_a = os.path.join(self.ultrapixel_path, stage_a) - self.stage_b = os.path.join(self.ultrapixel_path, stage_b) - self.stage_c = os.path.join(self.ultrapixel_path, stage_c) - self.effnet = os.path.join(self.ultrapixel_path, effnet) - self.previewer = os.path.join(self.ultrapixel_path, previewer) - self.controlnet = os.path.join(self.ultrapixel_path, controlnet) + self.stage_a = os.path.join(self.stablecascade_path, stage_a) + self.stage_b = os.path.join(self.stablecascade_path, stage_b) + self.stage_c = os.path.join(self.stablecascade_path, stage_c) + self.effnet = os.path.join(self.stablecascade_path, effnet) + self.previewer = os.path.join(self.stablecascade_path, previewer) + self.controlnet = os.path.join(self.stablecascade_path, controlnet) + self.previewer = os.path.join(self.stablecascade_path, previewer) + self.controlnet = os.path.join(self.stablecascade_path, controlnet) def set_config( self, @@ -113,8 +145,7 @@ class UltraPixel: captions = [self.prompt] height, width = self.height, self.width - pretrained_path = os.path.join(self.ultrapixel_path, self.pretrained) - sdd = torch.load(pretrained_path, map_location="cpu") + sdd = torch.load(self.pretrained, map_location="cpu") collect_sd = {} for k, v in sdd.items(): collect_sd[k[7:]] = v @@ -123,9 +154,8 @@ class UltraPixel: models.train_norm.load_state_dict(collect_sd) else: models.train_norm.load_state_dict(collect_sd, strict=True) - controlnet_path = os.path.join(self.ultrapixel_path, self.controlnet) models.controlnet.load_state_dict( - load_or_fail(controlnet_path), strict=True + load_or_fail(self.controlnet), strict=True ) models.generator.eval() diff --git a/workflow_controlnet.json b/workflow_controlnet.json index e004f23..148b468 100644 --- a/workflow_controlnet.json +++ b/workflow_controlnet.json @@ -1,6 +1,6 @@ { - "last_node_id": 14, - "last_link_id": 19, + "last_node_id": 15, + "last_link_id": 20, "nodes": [ { "id": 13, @@ -28,15 +28,79 @@ } }, { - "id": 1, + "id": 14, + "type": "UltraPixelProcess", + "pos": [ + 366, + 146 + ], + "size": { + "0": 450.48370361328125, + "1": 479.1488037109375 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "ULTRAPIXELMODEL", + "link": 20 + }, + { + "name": "controlnet_image", + "type": "IMAGE", + "link": 17 + } + ], + "outputs": [ + { + "name": "image", + "type": "IMAGE", + "links": [ + 18 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "edge_preview", + "type": "IMAGE", + "links": [ + 19 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "UltraPixelProcess" + }, + "widgets_values": [ + 2048, + 2048, + 973535413168831, + "randomize", + "bf16", + "true", + 10, + 1.1, + 20, + 4, + 0.7000000000000001, + "A close-up portrait of a young woman with flawless skin, vibrant red lipstick, and wavy brown hair, wearing a vintage floral dress and standing in front of a blooming garden, waving" + ] + }, + { + "id": 15, "type": "UltraPixelLoad", "pos": [ - 14, + 5, 139 ], "size": { "0": 315, - "1": 202 + "1": 250 }, "flags": {}, "order": 0, @@ -46,7 +110,7 @@ "name": "model", "type": "ULTRAPIXELMODEL", "links": [ - 16 + 20 ], "shape": 3, "slot_index": 0 @@ -62,15 +126,17 @@ "stage_c_bf16.safetensors", "effnet_encoder.safetensors", "previewer.safetensors", - "controlnet/canny.safetensors" + "controlnet/canny.safetensors", + "default", + "default" ] }, { "id": 5, "type": "LoadImage", "pos": [ - 16, - 395 + 2, + 452 ], "size": { "0": 315, @@ -104,70 +170,6 @@ "image" ] }, - { - "id": 14, - "type": "UltraPixelProcess", - "pos": [ - 366, - 146 - ], - "size": [ - 450.483717466484, - 479.1487988738271 - ], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [ - { - "name": "model", - "type": "ULTRAPIXELMODEL", - "link": 16 - }, - { - "name": "controlnet_image", - "type": "IMAGE", - "link": 17 - } - ], - "outputs": [ - { - "name": "image", - "type": "IMAGE", - "links": [ - 18 - ], - "shape": 3, - "slot_index": 0 - }, - { - "name": "edge_preview", - "type": "IMAGE", - "links": [ - 19 - ], - "shape": 3, - "slot_index": 1 - } - ], - "properties": { - "Node name for S&R": "UltraPixelProcess" - }, - "widgets_values": [ - 2048, - 2048, - 868772379065041, - "randomize", - "bf16", - "true", - 10, - 1.1, - 20, - 4, - 0.7000000000000001, - "A close-up portrait of a young woman with flawless skin, vibrant red lipstick, and wavy brown hair, wearing a vintage floral dress and standing in front of a blooming garden, waving" - ] - }, { "id": 12, "type": "PreviewImage", @@ -175,10 +177,10 @@ 1213, 140 ], - "size": [ - 1088.8128984964837, - 990.3649069338269 - ], + "size": { + "0": 1088.8128662109375, + "1": 990.3649291992188 + }, "flags": {}, "order": 3, "mode": 0, @@ -195,14 +197,6 @@ } ], "links": [ - [ - 16, - 1, - 0, - 14, - 0, - "ULTRAPIXELMODEL" - ], [ 17, 5, @@ -226,16 +220,24 @@ 13, 0, "IMAGE" + ], + [ + 20, + 15, + 0, + 14, + 0, + "ULTRAPIXELMODEL" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.683013455365071, + "scale": 0.5644739300537773, "offset": [ - 1.6029426679143006, - -24.445754383496105 + 129.91547966791447, + 91.38988261650393 ] } }, diff --git a/workflow_default.json b/workflow_default.json index 7b44fd2..b1a7f17 100644 --- a/workflow_default.json +++ b/workflow_default.json @@ -1,45 +1,7 @@ { - "last_node_id": 15, - "last_link_id": 19, + "last_node_id": 16, + "last_link_id": 20, "nodes": [ - { - "id": 1, - "type": "UltraPixelLoad", - "pos": [ - 14, - 139 - ], - "size": { - "0": 315, - "1": 202 - }, - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "model", - "type": "ULTRAPIXELMODEL", - "links": [ - 16 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "UltraPixelLoad" - }, - "widgets_values": [ - "ultrapixel_t2i.safetensors", - "stage_a.safetensors", - "stage_b_lite_bf16.safetensors", - "stage_c_bf16.safetensors", - "effnet_encoder.safetensors", - "previewer.safetensors", - "controlnet/canny.safetensors" - ] - }, { "id": 12, "type": "PreviewImage", @@ -83,7 +45,7 @@ { "name": "model", "type": "ULTRAPIXELMODEL", - "link": 16 + "link": 20 }, { "name": "controlnet_image", @@ -115,7 +77,7 @@ "widgets_values": [ 2048, 2048, - 741997913499515, + 940240851961031, "randomize", "bf16", "true", @@ -126,17 +88,49 @@ 0.7000000000000001, "A close-up portrait of a young woman with flawless skin, vibrant red lipstick, and wavy brown hair, wearing a vintage floral dress and standing in front of a blooming garden." ] + }, + { + "id": 16, + "type": "UltraPixelLoad", + "pos": [ + 7, + 148 + ], + "size": { + "0": 315, + "1": 250 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "model", + "type": "ULTRAPIXELMODEL", + "links": [ + 20 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "UltraPixelLoad" + }, + "widgets_values": [ + "ultrapixel_t2i.safetensors", + "stage_a.safetensors", + "stage_b_lite_bf16.safetensors", + "stage_c_bf16.safetensors", + "effnet_encoder.safetensors", + "previewer.safetensors", + "controlnet/canny.safetensors", + "default", + "default" + ] } ], "links": [ - [ - 16, - 1, - 0, - 14, - 0, - "ULTRAPIXELMODEL" - ], [ 18, 14, @@ -144,16 +138,24 @@ 12, 0, "IMAGE" + ], + [ + 20, + 16, + 0, + 14, + 0, + "ULTRAPIXELMODEL" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.6830134553650706, + "scale": 0.6830134553650707, "offset": [ - 9.066506738772325, - -54.63064767978551 + 34.98107673877221, + -18.467377679785468 ] } },