mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-07-31 18:27:08 +00:00
Moved ui preview-related classes out of io.py and into ui.py, refactored UIImages and related into PreviewImage and related
This commit is contained in:
Jedrzej Kosinski
parent
1ae7e7a1e2
commit
2999212480
@ -1207,7 +1207,7 @@ class NodeOutput:
|
||||
'''
|
||||
Standardized output of a node; can pass in any number of args and/or a UIOutput into 'ui' kwarg.
|
||||
'''
|
||||
def __init__(self, *args: Any, ui: UIOutput | dict=None, expand: dict=None, block_execution: str=None, **kwargs):
|
||||
def __init__(self, *args: Any, ui: _UIOutput | dict=None, expand: dict=None, block_execution: str=None, **kwargs):
|
||||
self.args = args
|
||||
self.ui = ui
|
||||
self.expand = expand
|
||||
@ -1219,21 +1219,7 @@ class NodeOutput:
|
||||
# TODO: use kwargs to refer to outputs by id + organize in proper order
|
||||
return self.args if len(self.args) > 0 else None
|
||||
|
||||
|
||||
class SavedResult:
|
||||
def __init__(self, filename: str, subfolder: str, type: FolderType):
|
||||
self.filename = filename
|
||||
self.subfolder = subfolder
|
||||
self.type = type
|
||||
|
||||
def as_dict(self):
|
||||
return {
|
||||
"filename": self.filename,
|
||||
"subfolder": self.subfolder,
|
||||
"type": self.type.value
|
||||
}
|
||||
|
||||
class UIOutput(ABC):
|
||||
class _UIOutput(ABC):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@ -1241,61 +1227,6 @@ class UIOutput(ABC):
|
||||
def as_dict(self) -> dict:
|
||||
... # TODO: finish
|
||||
|
||||
class UIImages(UIOutput):
|
||||
def __init__(self, values: list[SavedResult | dict], animated=False, **kwargs):
|
||||
self.values = values
|
||||
self.animated = animated
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"images": values,
|
||||
"animated": (self.animated,)
|
||||
}
|
||||
|
||||
class UILatents(UIOutput):
|
||||
def __init__(self, values: list[SavedResult | dict], **kwargs):
|
||||
self.values = values
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"latents": values,
|
||||
}
|
||||
|
||||
class UIAudio(UIOutput):
|
||||
def __init__(self, values: list[SavedResult | dict], **kwargs):
|
||||
self.values = values
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"audio": values,
|
||||
}
|
||||
|
||||
class UI3D(UIOutput):
|
||||
def __init__(self, values: list[SavedResult | dict], **kwargs):
|
||||
self.values = values
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"3d": values,
|
||||
}
|
||||
|
||||
class UIText(UIOutput):
|
||||
def __init__(self, value: str, **kwargs):
|
||||
self.value = value
|
||||
|
||||
def as_dict(self):
|
||||
return {"text": (self.value,)}
|
||||
|
||||
|
||||
def create_image_preview(image: Image.Type) -> UIImages:
|
||||
# TODO: finish, right now is just Cursor's hallucination
|
||||
return UIImages([SavedResult("preview.png", "comfy_org", FolderType.output)])
|
||||
|
||||
|
||||
class TestNode(ComfyNodeV3):
|
||||
@classmethod
|
||||
def DEFINE_SCHEMA(cls):
|
||||
|
||||
143
comfy_api/v3/ui.py
Normal file
143
comfy_api/v3/ui.py
Normal file
@ -0,0 +1,143 @@
|
||||
from __future__ import annotations
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from comfy_api.v3.io import Image, Mask, FolderType, _UIOutput
|
||||
# used for image preview
|
||||
import folder_paths
|
||||
import random
|
||||
from PIL import Image as PILImage
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
|
||||
class SavedResult:
|
||||
def __init__(self, filename: str, subfolder: str, type: FolderType):
|
||||
self.filename = filename
|
||||
self.subfolder = subfolder
|
||||
self.type = type
|
||||
|
||||
def as_dict(self):
|
||||
return {
|
||||
"filename": self.filename,
|
||||
"subfolder": self.subfolder,
|
||||
"type": self.type
|
||||
}
|
||||
|
||||
class PreviewImage(_UIOutput):
|
||||
def __init__(self, image: Image.Type, animated: bool=False, **kwargs):
|
||||
output_dir = folder_paths.get_temp_directory()
|
||||
type = "temp"
|
||||
prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
||||
compress_level = 1
|
||||
filename_prefix = "ComfyUI"
|
||||
|
||||
filename_prefix += prefix_append
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir, image[0].shape[1], image[0].shape[0])
|
||||
results = list()
|
||||
for (batch_number, image) in enumerate(image):
|
||||
i = 255. * image.cpu().numpy()
|
||||
img = PILImage.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
||||
metadata = None
|
||||
# if not args.disable_metadata:
|
||||
# metadata = PngInfo()
|
||||
# if prompt is not None:
|
||||
# metadata.add_text("prompt", json.dumps(prompt))
|
||||
# if extra_pnginfo is not None:
|
||||
# for x in extra_pnginfo:
|
||||
# metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
|
||||
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
||||
file = f"{filename_with_batch_num}_{counter:05}_.png"
|
||||
img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=compress_level)
|
||||
results.append(SavedResult(file, subfolder, type))
|
||||
counter += 1
|
||||
|
||||
self.values = results
|
||||
self.animated = animated
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"images": values,
|
||||
"animated": (self.animated,)
|
||||
}
|
||||
|
||||
class PreviewMask(PreviewImage):
|
||||
def __init__(self, mask: PreviewMask.Type, animated: bool=False, **kwargs):
|
||||
preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
||||
super().__init__(preview, animated, **kwargs)
|
||||
|
||||
# class UILatent(_UIOutput):
|
||||
# def __init__(self, values: list[SavedResult | dict], **kwargs):
|
||||
# output_dir = folder_paths.get_temp_directory()
|
||||
# type = "temp"
|
||||
# prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
||||
# compress_level = 1
|
||||
# filename_prefix = "ComfyUI"
|
||||
|
||||
|
||||
# full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
|
||||
|
||||
# # support save metadata for latent sharing
|
||||
# prompt_info = ""
|
||||
# if prompt is not None:
|
||||
# prompt_info = json.dumps(prompt)
|
||||
|
||||
# metadata = None
|
||||
# if not args.disable_metadata:
|
||||
# metadata = {"prompt": prompt_info}
|
||||
# if extra_pnginfo is not None:
|
||||
# for x in extra_pnginfo:
|
||||
# metadata[x] = json.dumps(extra_pnginfo[x])
|
||||
|
||||
# file = f"{filename}_{counter:05}_.latent"
|
||||
|
||||
# results: list[FileLocator] = []
|
||||
# results.append({
|
||||
# "filename": file,
|
||||
# "subfolder": subfolder,
|
||||
# "type": "output"
|
||||
# })
|
||||
|
||||
# file = os.path.join(full_output_folder, file)
|
||||
|
||||
# output = {}
|
||||
# output["latent_tensor"] = samples["samples"].contiguous()
|
||||
# output["latent_format_version_0"] = torch.tensor([])
|
||||
|
||||
# comfy.utils.save_torch_file(output, file, metadata=metadata)
|
||||
|
||||
# self.values = values
|
||||
|
||||
# def as_dict(self):
|
||||
# values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
# return {
|
||||
# "latents": values,
|
||||
# }
|
||||
|
||||
class PreviewAudio(_UIOutput):
|
||||
def __init__(self, values: list[SavedResult | dict], **kwargs):
|
||||
self.values = values
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"audio": values,
|
||||
}
|
||||
|
||||
class PreviewUI3D(_UIOutput):
|
||||
def __init__(self, values: list[SavedResult | dict], **kwargs):
|
||||
self.values = values
|
||||
|
||||
def as_dict(self):
|
||||
values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
|
||||
return {
|
||||
"3d": values,
|
||||
}
|
||||
|
||||
class PreviewText(_UIOutput):
|
||||
def __init__(self, value: str, **kwargs):
|
||||
self.value = value
|
||||
|
||||
def as_dict(self):
|
||||
return {"text": (self.value,)}
|
||||
@ -1,5 +1,5 @@
|
||||
import torch
|
||||
from comfy_api.v3 import io
|
||||
from comfy_api.v3 import io, ui
|
||||
import logging
|
||||
import folder_paths
|
||||
import comfy.utils
|
||||
@ -96,7 +96,7 @@ class V3TestNode(io.ComfyNodeV3):
|
||||
if hasattr(cls, "doohickey"):
|
||||
raise Exception("The 'cls' variable leaked state on class properties between runs!")
|
||||
cls.doohickey = "LOLJK"
|
||||
return io.NodeOutput(some_int, image)
|
||||
return io.NodeOutput(some_int, image, ui=ui.PreviewImage(image))
|
||||
|
||||
|
||||
class V3LoraLoader(io.ComfyNodeV3):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user