# !/usr/bin/env python # -*- coding: UTF-8 -*- import numpy as np import torch import os from .Chrono.scripts.run_inference_diffusers import load_dit_model,load_prompt_enhancer,prompt_enhance,inference_chrono,load_prompt_enhancer_cf from .model_loader_utils import tensor_upscale,tensor2image import folder_paths from typing_extensions import override from comfy_api.latest import ComfyExtension, io import nodes from pathlib import PureWindowsPath from transformers import AutoProcessor import comfy.model_management as mm MAX_SEED = np.iinfo(np.int32).max node_cr_path = os.path.dirname(os.path.abspath(__file__)) device = torch.device( "cuda:0") if torch.cuda.is_available() else torch.device( "mps") if torch.backends.mps.is_available() else torch.device( "cpu") weigths_gguf_current_path = os.path.join(folder_paths.models_dir, "gguf") if not os.path.exists(weigths_gguf_current_path): os.makedirs(weigths_gguf_current_path) folder_paths.add_model_folder_path("gguf", weigths_gguf_current_path) # gguf dir class ChronoEdit_SM_Model(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_Model", display_name="ChronoEdit_SM_Model", category="ChronoEdit_SM", inputs=[ io.Combo.Input("dit",options= ["none"] + folder_paths.get_filename_list("diffusion_models") ), io.Combo.Input("gguf",options= ["none"] + folder_paths.get_filename_list("gguf")), ], outputs=[ io.Custom("ChronoEdit_SM_Model").Output(display_name="model"), ], ) @classmethod def execute(cls, dit,gguf) -> io.NodeOutput: dit_path=folder_paths.get_full_path("diffusion_models", dit) if dit != "none" else None gguf_path=folder_paths.get_full_path("gguf", gguf) if gguf != "none" else None model=load_dit_model(dit_path, gguf_path,node_cr_path) return io.NodeOutput(model) class ChronoEdit_SM_Lora(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_Lora", display_name="ChronoEdit_SM_Lora", category="ChronoEdit_SM", inputs=[ io.Custom("ChronoEdit_SM_Model").Input("model"), io.Combo.Input("lora",options= ["none"] + folder_paths.get_filename_list("loras") ), io.Float.Input("lora_scale", default=1, min=0.1, max=1.0,step=0.1, round=0.01,), ], outputs=[ io.Custom("ChronoEdit_SM_Model").Output(display_name="model"), ], ) @classmethod def execute(cls, model,lora,lora_scale) -> io.NodeOutput: lora_path=folder_paths.get_full_path("loras", lora) if lora != "none" else None if lora_path is not None: print(f"Loading LoRA weights from {lora_path}...") all_adapters = model.get_list_adapters() dit_list=[] if all_adapters: dit_list= all_adapters['transformer'] if dit_list: if os.path.splitext(os.path.basename(lora_path))[0] in dit_list: pass else: model.load_lora_weights(lora_path, adapter_name=os.path.splitext(os.path.basename(lora_path))[0]) else: model.load_lora_weights(lora_path,adapter_name= os.path.splitext(os.path.basename(lora_path))[0]) else: model.load_lora_weights(lora_path,adapter_name= os.path.splitext(os.path.basename(lora_path))[0]) model.set_adapters([os.path.splitext(os.path.basename(lora_path))[0]], adapter_weights=[lora_scale]) try: dit_list= model.get_list_adapters()['transformer'] for name in dit_list: if lora_path is not None: if name!=os.path.splitext(os.path.basename(lora_path))[0]: model.delete_adapters(name) print(f"当前激活的适配器: {os.path.splitext(os.path.basename(lora_path))[0]}") else: model.delete_adapters(name) except:pass return io.NodeOutput(model) class ChronoEdit_SM_Vae(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_Vae", display_name="ChronoEdit_SM_Vae", category="ChronoEdit_SM", inputs=[ io.Vae.Input("vae"), io.Latent.Input("latent"), io.Combo.Input("vae_decoder",options= ["none"] + folder_paths.get_filename_list("vae") ), ], outputs=[ io.Image.Output(display_name="image"), ], ) @classmethod def execute(cls, vae,latent,vae_decoder) -> io.NodeOutput: samples=latent["samples"] if samples.shape[2]>2: #torch.Size([1, 16, 8, 60, 104]) if vae_decoder!="none": video_edit = vae.decode(samples[:, :, [0, -1]]) video_reason =vae.decode(samples[:, :, :-1]) # print(f"Video edit shape: {video_edit.shape}") #Video edit shape: torch.Size([1, 5, 480, 832, 3]) # print(f"video_reason edit shape: {video_reason.shape}") #Video edit shape: torch.Size([1, 25, 480, 832, 3]) #images = torch.cat([video_reason, video_edit[:, :, 1:]], dim=2) images = torch.cat([video_reason, video_edit[:, 1:, :, :]], dim=1) else: images = vae.decode(samples) else: images = vae.decode(samples) if len(images.shape) == 5: #Combine batches images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) return io.NodeOutput(images) class ChronoEdit_SM_Enhance_Loader(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_Enhance_Loader", display_name="ChronoEdit_SM_Enhance_Loader", category="ChronoEdit_SM", inputs=[ io.String.Input("repo", multiline=False, default="Qwen/Qwen2.5-VL-7B-Instruct"), io.Combo.Input("clip",options= ["none"] + folder_paths.get_filename_list("clip")), ], outputs=[ io.Custom("ChronoEdit_SM_Model_En").Output(display_name="model"), ], ) @classmethod def execute(cls, repo,clip) -> io.NodeOutput: clip_path=folder_paths.get_full_path("clip", clip) if clip != "none" else None if clip_path is not None: if "2.5" in clip_path.lower() or "qwen2" in clip_path.lower(): model_name = os.path.join(node_cr_path,"Qwen/Qwen2.5-VL-7B-Instruct") elif "qwen3" in clip_path.lower(): model_name = os.path.join(node_cr_path,"Qwen/Qwen3-VL-8B-Instruct") else: raise ValueError(f"Unsupported model: {clip_path}") model=load_prompt_enhancer_cf(clip_path,model_name) processor = AutoProcessor.from_pretrained(model_name) elif repo: repo=PureWindowsPath(repo) model,processor=load_prompt_enhancer(repo) else: raise ValueError(f"Unsupported model: {repo} or {clip}") pipe={"model":model,"processor":processor,} return io.NodeOutput(pipe) class ChronoEdit_SM_Enhance(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_Enhance", display_name="ChronoEdit_SM_Enhance", category="ChronoEdit_SM", inputs=[ io.Image.Input("image"), io.String.Input("prompt", multiline=True, default="Transform the image so that inside the floral teacup of steaming tea, a small, cute mouse is sitting and taking a bath; the mouse should look relaxed and cheerful, with a tiny white bath towel draped over its head as if enjoying a spa moment, while the steam rises gently around it, blending seamlessly with the warm and cozy atmosphere."), io.Custom("ChronoEdit_SM_Model_En").Input("model"), ], outputs=[ io.String.Output(display_name="prompt"), ], ) @classmethod def execute(cls, image,prompt,model) -> io.NodeOutput: image=tensor2image(image) prompt=prompt_enhance(image,prompt,model["model"],model["processor"]) return io.NodeOutput(prompt) class ChronoEdit_SM_Latent(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_Latent", display_name="ChronoEdit_SM_Latent", category="ChronoEdit_SM", inputs=[ io.Vae.Input("vae"), io.Image.Input("image"), io.Int.Input("width", default=832, min=128, max=nodes.MAX_RESOLUTION,step=16,display_mode=io.NumberDisplay.number), io.Int.Input("height", default=480, min=128, max=nodes.MAX_RESOLUTION,step=16,display_mode=io.NumberDisplay.number), io.Combo.Input("num_frames",options= [5,29]), ], outputs=[ io.Latent.Output(display_name="Latent"), io.Int.Output(display_name="num_frames"), ], ) @classmethod def execute(cls, vae,image,width,height,num_frames) -> io.NodeOutput: image = tensor_upscale(image, width, height) image = torch.cat([image, image.new_zeros(int(num_frames) - 1, image.shape[1], image.shape[2], image.shape[3])], dim=0) Latent=vae.encode(image) return io.NodeOutput(Latent,num_frames) class ChronoEdit_SM_KSampler(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="ChronoEdit_SM_KSampler", display_name="ChronoEdit_SM_KSampler", category="ChronoEdit_SM", inputs=[ io.Custom("ChronoEdit_SM_Model").Input("model"), io.Latent.Input("image_latent"), #1, 16, 2, 60, 104 io.Conditioning.Input("positive"), io.Conditioning.Input("negative"), io.Int.Input("num_frames",force_input=True), io.Int.Input("seed", default=0, min=0, max=MAX_SEED), io.Float.Input("flow_shift", default=2.0, min=0.1, max=10.0, step=0.1, round=0.01,), io.Int.Input("steps", default=8, min=1, max=10000), io.Float.Input("guidance_scale", default=1.0, min=0.1, max=10.0, step=0.1, round=0.01,), io.Int.Input("num_temporal_reasoning_steps", default=50, min=0, max=MAX_SEED,step=1,), io.Boolean.Input("offload_model", default=True), io.Int.Input("block_num", default=1, min=1, max=48,step=1), ], outputs=[ io.Latent.Output(display_name="Latent"), ], ) @classmethod def execute(cls, model,image_latent,positive,negative,num_frames,seed,flow_shift, steps, guidance_scale,num_temporal_reasoning_steps,offload_model,block_num) -> io.NodeOutput: cf_models=mm.loaded_models() try: for pipe in cf_models: pipe.unpatch_model(device_to=torch.device("cpu")) print(f"Unpatching models.{pipe}") except: pass mm.soft_empty_cache() torch.cuda.empty_cache() max_gpu_memory = torch.cuda.max_memory_allocated() print(f"After Max GPU memory allocated: {max_gpu_memory / 1000 ** 3:.2f} GB") width,height=image_latent.shape[3]*8,image_latent.shape[4]*8 Latent=inference_chrono (model,image_latent,positive[0][0],positive[0][1]["unclip_conditioning"][0]["clip_vision_output"]["penultimate_hidden_states"],negative[0][0],flow_shift,seed,steps, width,height,num_frames,num_temporal_reasoning_steps,guidance_scale,offload_model,block_num,device) out={} out["samples"]=Latent #torch.Size([1, 16, 2, 60, 104]) return io.NodeOutput(out) from aiohttp import web from server import PromptServer @PromptServer.instance.routes.get("/ChronoEdit_SM_Extension") async def get_hello(request): return web.json_response("ChronoEdit_SM_Extension") class ChronoEdit_SM_Extension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ ChronoEdit_SM_Model, ChronoEdit_SM_Lora, ChronoEdit_SM_Enhance_Loader, ChronoEdit_SM_Enhance, ChronoEdit_SM_Vae, ChronoEdit_SM_Latent, ChronoEdit_SM_KSampler, ] async def comfy_entrypoint() -> ChronoEdit_SM_Extension: # ComfyUI calls this to load your extension and its nodes. return ChronoEdit_SM_Extension()