Add custom dirs

This commit is contained in:
2kpr
2024-07-17 12:49:36 -05:00
parent 3695d90939
commit 3ea5ef9865
5 changed files with 248 additions and 158 deletions
+2
View File
@@ -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."
<br/>
+59 -5
View File
@@ -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,)
+42 -12
View File
@@ -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()
+90 -88
View File
@@ -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
]
}
},
+55 -53
View File
@@ -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
]
}
},