nodes/TashaSkyUp-EternalKernelPytorchNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

tinytxt2imgv2

Browse Source
This commit is contained in:
TashaSkyUp
2025-01-09 19:33:25 -07:00
parent 4aa9c36b7c
commit 43ffaaf78a
1 changed files with 259 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image.py
+259
View File
@@ -288,6 +288,7 @@ class TinyTxtToImg:
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
"one_seed": ([True, False], {"default": False}),
"unload_model": ([True, False], {"default": True}),
"func_only": ([True, False], {"default": True}),
},
"optional": {
@@ -507,6 +508,264 @@ class TinyTxtToImg:
gc.collect()
return ret
@ETK_image_base
class TinyTxtToImgV2:
"""small text to image generator"""
share_clip = None
share_mdl = None
share_vae = None
def __init__(self):
print("ETK> TinyTxtToImgV2 init")
self.mdl = None
self.clp = None
self.vae = None
self.vision = None
self.steps = 10
self.cfg = 8
self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0]
self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
self.positive = None
self.negative = None
self.width = 512
self.height = 512
self.batch_size = 1
self.denoise = 1.0
self.latent_image = None
self.samples = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True}),
"neg_prompt": ("STRING", {"multiline": True}),
"clip_encoder": (["comfy -ignore below", "advanced"],),
"token_normalization": (["none", "mean", "length", "length+mean"],),
"weight_interpretation": (["comfy", "A1111", "compel", "comfy++"],),
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
"one_seed": ([True, False], {"default": False}),
"unload_model": ([True, False], {"default": True}),
"func_only": ([True, False], {"default": True}),
},
"optional": {
"name": ("STRING", {"multiline": False, "default": "tinytxt2img"}),
"overrides": ("STRING", {"multiline": True, "default": '{"width":768,"height":768,"steps":20}'}),
"latent": ("LATENT",),
"pos_cond": ("CONDITIONING",),
"neg_cond": ("CONDITIONING",),
"image": ("IMAGE",),
}
}
CATEGORY = "ETK"
RETURN_TYPES = ("IMAGE", "FUNC", "LATENT", "CONDITIONING", "CONDITIONING",)
RETURN_NAMES = ("image", "FUNC", "latent", "positive", "negative",)
FUNCTION = "tinytxt2img"
def tinytxt2img(self, **kwargs):
"""If func_only is True, return a callable function instead of immediately generating an image."""
import torch
import comfy.model_management as mm
import random
import json
import gc
from nodes import CLIPTextEncode
torch.cuda.empty_cache()
def _do_generation(**override_kwargs):
"""This inner function carries out the entire generation logic, using merged arguments."""
# Merge original kwargs with any overrides passed to _do_generation
all_args = dict(kwargs)
all_args.update(override_kwargs)
unload_model = all_args.get("unload_model", True)
prompt = all_args.get("prompt", "")
neg_prompt = all_args.get("neg_prompt", "")
clip_encoder = all_args.get("clip_encoder", "comfy -ignore below")
token_normalization = all_args.get("token_normalization", "length+mean")
weight_interpretation = all_args.get("weight_interpretation", "comfy++")
ckpt_name = all_args.get("ckpt_name", None)
model_ref = all_args.get("model", None)
vae = all_args.get("vae", None)
clip = all_args.get("clip", None)
overrides = all_args.get("overrides", "")
name = all_args.get("name", "tinytxt2img")
FUNC = all_args.get("FUNC", None)
render_what = all_args.get("render_what", "image")
self.latent_image = all_args.get("latent", None)
self.pos_cond = all_args.get("pos_cond", None)
self.neg_cond = all_args.get("neg_cond", None)
self.seed = all_args.get("seed", random.randint(0, 2 ** 32 - 1))
self.one_seed_per_batch = all_args.get("one_seed", False)
self.denoise = all_args.get("denoise", 1.0)
self.steps = all_args.get("steps", 10)
if model_ref and vae and clip:
self.mdl = model_ref
self.vae = vae
self.clp = clip
else:
from nodes import CheckpointLoaderSimple
self.mdl, self.clp, self.vae = CheckpointLoaderSimple.load_checkpoint(
None, ckpt_name=ckpt_name
)
# Check if an input image is provided
input_image = all_args.get("image", None)
if input_image is not None:
with torch.no_grad():
self.vae.device = torch.device("cuda:0")
self.vae.first_stage_model = self.vae.first_stage_model.to("cuda:0")
input_image = input_image.to("cuda")
q = self.vae.encode(input_image)
self.latent_image = {"samples": q}
self.cfg = all_args.get("cfg", 8)
# Sampler logic
if "sampler_name" not in all_args and ckpt_name:
if "lcm" in ckpt_name.lower():
self.sampler_name = comfy.samplers.KSampler.SAMPLERS[18]
self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
else:
self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0]
self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
else:
self.sampler_name = all_args.get("sampler_name", comfy.samplers.KSampler.SAMPLERS[0])
self.scheduler = all_args.get("scheduler", comfy.samplers.KSampler.SCHEDULERS[0])
self.positive = prompt
self.negative = neg_prompt
self.width = 512
self.height = 512
self.batch_size = 1
if not self.denoise:
self.denoise = 1.0
# Process overrides if provided
if overrides:
if overrides != "":
overrides = overrides.replace("'", "\"")
if overrides[0] != "{":
overrides = "{" + overrides + "}"
valid_dict = json.loads(overrides)
for k, v in valid_dict.items():
if k in self.__dict__:
self.__setattr__(k, v)
else:
print(f"invalid override key: {k}")
if FUNC is not None:
self.__dict__.update(FUNC(self.__dict__))
# Build positive/negative conditioning if not provided
if not self.pos_cond:
# Possibly advanced clip encode
self.clp.cond_stage_model = self.clp.cond_stage_model.to("cuda")
if clip_encoder == "comfy -ignore below":
self.clp_encode = lambda x: CLIPTextEncode.encode(None, self.clp, x)[0]
self.pos_cond = self.clp_encode(self.positive)
self.neg_cond = self.clp_encode(self.negative)
else:
from nodes import CLIPTextEncodeAdvanced
self.clp_encode = lambda x: CLIPTextEncodeAdvanced.encode(
None, self.clp, x, token_normalization, weight_interpretation
)[0]
try:
self.pos_cond = self.clp_encode(self.positive)
self.neg_cond = self.clp_encode(self.negative)
except Exception as e:
print(e)
raise ValueError("advanced clip encoder failed")
self.clp.cond_stage_model = self.clp.cond_stage_model.to("cpu")
torch.cuda.empty_cache()
gc.collect()
if self.latent_image is None:
from nodes import EmptyLatentImage
self.latent_image = EmptyLatentImage().generate(
self.width, self.height, self.batch_size
)[0]
image = None
# Only sample/decode if the user wants "latent" or "image" or "all"
if "latent" in render_what or "all" in render_what or "image" in render_what:
self.KSampler = ETKKSampler()
self.samples = self.KSampler.sample(
model=self.mdl,
seed=self.seed,
steps=self.steps,
cfg=self.cfg,
sampler_name=self.sampler_name,
scheduler=self.scheduler,
latent_image=self.latent_image,
pos_cond=self.pos_cond,
neg_cond=self.neg_cond,
positive=self.pos_cond,
negative=self.neg_cond,
denoise=self.denoise,
one_seed_per_batch=self.one_seed_per_batch
)[0]
if "image" in render_what or "all" in render_what:
self.samples["samples"] = self.samples["samples"].to("cuda")
self.vae.first_stage_model = self.vae.first_stage_model.to("cuda:0")
self.vae.device = torch.device("cuda:0")
torch.cuda.empty_cache()
decoded_images = self.vae.decode(self.samples["samples"])
image = decoded_images.to("cpu")
self.vae.first_stage_model.to("cpu")
self.vae.device = torch.device("cpu")
torch.cuda.empty_cache()
gc.collect()
self.FUNC = self.tinytxt2img
self.ARGS = all_args
self.samples["samples"] = self.samples["samples"].to("cpu")
ret = (
image,
self,
self.samples,
self.pos_cond,
self.neg_cond
)
# If requested, unload
if unload_model:
del self.mdl
del self.clp
del self.vae
if "vision" in self.__dict__:
del self.vision
del self.samples
del self.FUNC
del self.ARGS
del self.pos_cond
del self.neg_cond
del self.KSampler
mm.free_memory(mm.get_total_memory(mm.get_torch_device()), mm.get_torch_device())
torch.cuda.empty_cache()
gc.collect()
return ret
# If func_only is True, return a placeholder 5-tuple where the second item is our generator
if kwargs.get("func_only", False):
return (None, _do_generation, None, None, None)
else:
# Otherwise, perform the entire generation immediately and return the standard 5-tuple
return _do_generation()
@ETK_image_base
class PromptTemplate: