149 lines
4.5 KiB
Python
149 lines
4.5 KiB
Python
import os
|
|
import folder_paths
|
|
import torch
|
|
from . import ultrapixel as up
|
|
from huggingface_hub import snapshot_download
|
|
|
|
|
|
def get_models(base):
|
|
models_path = os.path.join(folder_paths.models_dir, "ultrapixel")
|
|
models = [base]
|
|
if os.path.exists(models_path):
|
|
for path in os.listdir(models_path):
|
|
if path.lower().endswith(".safetensors"):
|
|
if path not in models:
|
|
models.append(path)
|
|
return models
|
|
|
|
|
|
def download_models():
|
|
models = [
|
|
["ultrapixel_t2i.safetensors", "2kpr/UltraPixel"],
|
|
["stage_a.safetensors", "stabilityai/stable-cascade"],
|
|
["previewer.safetensors", "stabilityai/stable-cascade"],
|
|
["effnet_encoder.safetensors", "stabilityai/stable-cascade"],
|
|
["stage_b_lite_bf16.safetensors ", "stabilityai/stable-cascade"],
|
|
["stage_c_bf16.safetensors", "stabilityai/stable-cascade"],
|
|
]
|
|
for model in models:
|
|
model_path = os.path.join(folder_paths.models_dir, "ultrapixel")
|
|
if os.path.exists(os.path.join(model_path, model[0])):
|
|
# print("Path already exists:", os.path.join(model_path, model[0]))
|
|
continue
|
|
print(f"Downloading from {model[1]} to:", os.path.join(model_path, model[0]))
|
|
snapshot_download(
|
|
repo_id=model[1],
|
|
allow_patterns=[model[0]],
|
|
local_dir=model_path,
|
|
)
|
|
|
|
|
|
class UltraPixelLoad:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ultrapixel": (get_models("ultrapixel_t2i.safetensors"),),
|
|
"stage_a": (get_models("stage_a.safetensors"),),
|
|
"stage_b": (get_models("stage_b_lite_bf16.safetensors"),),
|
|
"stage_c": (get_models("stage_c_bf16.safetensors"),),
|
|
"effnet": (get_models("effnet_encoder.safetensors"),),
|
|
"previewer": (get_models("previewer.safetensors"),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("ULTRAPIXELMODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "ultrapixel"
|
|
CATEGORY = "UltraPixel"
|
|
|
|
def ultrapixel(self, ultrapixel, stage_a, stage_b, stage_c, effnet, previewer):
|
|
download_models()
|
|
model = up.UltraPixel(ultrapixel, stage_a, stage_b, stage_c, effnet, previewer)
|
|
return (model,)
|
|
|
|
|
|
class UltraPixelProcess:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("ULTRAPIXELMODEL",),
|
|
"height": (
|
|
"INT",
|
|
{"default": 2048, "min": 512, "max": 5120, "step": 256},
|
|
),
|
|
"width": (
|
|
"INT",
|
|
{"default": 2048, "min": 512, "max": 5120, "step": 256},
|
|
),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"dtype": (["bf16", "fp32"],),
|
|
"stage_a_tiled": (["true", "false"],),
|
|
"stage_b_steps": (
|
|
"INT",
|
|
{"default": 10, "min": 1, "max": 1000, "step": 1},
|
|
),
|
|
"stage_b_cfg": (
|
|
"FLOAT",
|
|
{"default": 1.1, "min": 1.0, "max": 1000.0, "step": 0.1},
|
|
),
|
|
"stage_c_steps": (
|
|
"INT",
|
|
{"default": 20, "min": 1, "max": 1000, "step": 1},
|
|
),
|
|
"stage_c_cfg": (
|
|
"FLOAT",
|
|
{"default": 4.0, "min": 1.0, "max": 1000.0, "step": 0.1},
|
|
),
|
|
"prompt": (
|
|
"STRING",
|
|
{"multiline": True, "dynamicPrompts": True},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "ultrapixel"
|
|
CATEGORY = "UltraPixel"
|
|
|
|
def ultrapixel(
|
|
self,
|
|
model,
|
|
height,
|
|
width,
|
|
seed,
|
|
dtype,
|
|
stage_a_tiled,
|
|
stage_b_steps,
|
|
stage_b_cfg,
|
|
stage_c_steps,
|
|
stage_c_cfg,
|
|
prompt,
|
|
):
|
|
model.set_config(
|
|
height,
|
|
width,
|
|
seed,
|
|
dtype,
|
|
stage_a_tiled,
|
|
stage_b_steps,
|
|
stage_b_cfg,
|
|
stage_c_steps,
|
|
stage_c_cfg,
|
|
prompt,
|
|
)
|
|
image = model.process()
|
|
return (image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UltraPixelLoad": UltraPixelLoad,
|
|
"UltraPixelProcess": UltraPixelProcess,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UltraPixelLoad": "UltraPixel Load",
|
|
"UltraPixelProcess": "UltraPixel Process",
|
|
}
|