Upload run_generator.py

run_generator.py

@@ -0,0 +1,234 @@
+# 이 파일은 셀 3에서 서브프로세스로 실행됩니다.
+
+import sys
+import os
+import time
+import glob
+import gc
+import torch
+import subprocess
+import random
+import argparse
+from typing import Sequence, Mapping, Any, Union
+
+# --- 0. 기본 설정 및 인수 파싱 --
+def parse_args():
+    parser = argparse.ArgumentParser(description="ComfyUI Video Generation Script")
+    parser.add_argument("--positive_prompt", type=str, required=True)
+    parser.add_argument("--negative_prompt", type=str, required=True)
+    parser.add_argument("--width", type=int, required=True)
+    parser.add_argument("--height", type=int, required=True)
+    parser.add_argument("--length", type=int, required=True)
+    parser.add_argument("--upscale_ratio", type=float, required=True)
+    parser.add_argument("--custom_lora_1_name", type=str, default="None")
+    parser.add_argument("--custom_lora_1_strength_model", type=float, default=1.0)
+    parser.add_argument("--custom_lora_1_strength_clip", type=float, default=1.0)
+    parser.add_argument("--custom_lora_2_name", type=str, default="None")
+    parser.add_argument("--custom_lora_2_strength_model", type=float, default=1.0)
+    parser.add_argument("--custom_lora_2_strength_clip", type=float, default=1.0)
+    return parser.parse_args()
+
+def clear_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    gc.collect()
+
+# --- 1. ComfyUI 관련 함수 정의 ---
+COMFYUI_BASE_PATH = '/content/ComfyUI'
+
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    try:
+        return obj[index]
+    except (KeyError, TypeError):
+        if isinstance(obj, dict) and "result" in obj:
+            return obj["result"][index]
+        raise
+
+def add_comfyui_directory_to_sys_path() -> None:
+    if os.path.isdir(COMFYUI_BASE_PATH) and COMFYUI_BASE_PATH not in sys.path:
+        sys.path.append(COMFYUI_BASE_PATH)
+        print(f"'{COMFYUI_BASE_PATH}' added to sys.path")
+
+def import_custom_nodes() -> None:
+    try:
+        import nest_asyncio
+        nest_asyncio.apply()
+    except ImportError:
+        subprocess.run([sys.executable, "-m", "pip", "install", "-q", "nest_asyncio"])
+        import nest_asyncio
+        nest_asyncio.apply()
+
+    import asyncio
+    import execution
+    from nodes import init_extra_nodes
+    import server
+
+    loop = asyncio.get_event_loop() or asyncio.new_event_loop()
+    server_instance = server.PromptServer(loop)
+    execution.PromptQueue(server_instance)
+    if not loop.is_running():
+        loop.run_until_complete(init_extra_nodes())
+    else:
+        asyncio.ensure_future(init_extra_nodes())
+
+# --- 2. 메인 실행 로직 ---
+def main():
+    args = parse_args()
+
+    print("🚀 동영상 생성을 시작합니다...")
+    print(f"프롬프트: {args.positive_prompt[:50]}...")
+    print(f"크기: {args.width}x{args.height}, 길이: {args.length} 프레임")
+    print(f"최종 업스케일 비율: {args.upscale_ratio}x")
+    if args.custom_lora_1_name != "None":
+        print(f"커스텀 LoRA 1: {args.custom_lora_1_name} (강도: {args.custom_lora_1_strength_model})")
+    if args.custom_lora_2_name != "None":
+        print(f"커스텀 LoRA 2: {args.custom_lora_2_name} (강도: {args.custom_lora_2_strength_model})")
+
+
+    subprocess.run(f"rm -rf {COMFYUI_BASE_PATH}/output/up/*", shell=True, check=True)
+    subprocess.run(f"rm -rf {COMFYUI_BASE_PATH}/output/temp/*", shell=True, check=True)
+
+    add_comfyui_directory_to_sys_path()
+
+    from utils.extra_config import load_extra_path_config
+    extra_model_paths_file = os.path.join(COMFYUI_BASE_PATH, "extra_model_paths.yaml")
+    if os.path.exists(extra_model_paths_file):
+        load_extra_path_config(extra_model_paths_file)
+
+    print("Initializing ComfyUI custom nodes...")
+    import_custom_nodes()
+    from nodes import NODE_CLASS_MAPPINGS
+    print("Custom nodes initialized successfully.")
+
+    with torch.inference_mode():
+        unetloadergguf = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]()
+        cliploader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
+        loraloader = NODE_CLASS_MAPPINGS["LoraLoader"]()
+        cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+        vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+        loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+        upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
+        mxslider = NODE_CLASS_MAPPINGS["mxSlider"]()
+        vhs_loadimagespath = NODE_CLASS_MAPPINGS["VHS_LoadImagesPath"]()
+        modelsamplingsd3 = NODE_CLASS_MAPPINGS["ModelSamplingSD3"]()
+        wanimagetovideo = NODE_CLASS_MAPPINGS["WanImageToVideo"]()
+        ksampleradvanced = NODE_CLASS_MAPPINGS["KSamplerAdvanced"]()
+        vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
+        rife_vfi = NODE_CLASS_MAPPINGS["RIFE VFI"]()
+        vhs_videocombine = NODE_CLASS_MAPPINGS["VHS_VideoCombine"]()
+        easy_mathfloat = NODE_CLASS_MAPPINGS["easy mathFloat"]()
+        imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
+        imagescaleby = NODE_CLASS_MAPPINGS["ImageScaleBy"]()
+        saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
+
+
+        print("Starting Step 1: Initial Sampling")
+        cliploader_38 = cliploader.load_clip(clip_name="umt5_xxl_fp8_e4m3fn_scaled.safetensors", type="wan", device="default")
+        unetloadergguf_84 = unetloadergguf.load_unet(unet_name="Wan2.2-I2V-A14B-HighNoise-Q4_K_S.gguf")
+
+        # 1번 샘플러 LoRA 체인
+        model, clip = loraloader.load_lora(lora_name="lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank32_bf16.safetensors", strength_model=2, strength_clip=2, model=get_value_at_index(unetloadergguf_84, 0), clip=get_value_at_index(cliploader_38, 0))
+        if args.custom_lora_1_name and args.custom_lora_1_name != "None":
+            print(f"1번 샘플러에 커스텀 LoRA 적용: {args.custom_lora_1_name}")
+            model, clip = loraloader.load_lora(lora_name=args.custom_lora_1_name, strength_model=args.custom_lora_1_strength_model, strength_clip=args.custom_lora_1_strength_clip, model=model, clip=clip)
+        loraloader_78 = loraloader.load_lora(lora_name="FastWan_T2V_14B_480p_lora_rank_128_bf16.safetensors", strength_model=1.5, strength_clip=1.5, model=model, clip=clip)
+
+        cliptextencode_6 = cliptextencode.encode(text=args.positive_prompt, clip=get_value_at_index(loraloader_78, 1))
+        cliptextencode_7 = cliptextencode.encode(text=args.negative_prompt, clip=get_value_at_index(loraloader_78, 1))
+        vaeloader_39 = vaeloader.load_vae(vae_name="wan_2.1_vae.safetensors")
+        loadimage_62 = loadimage.load_image(image="example.png")
+        wanimagetovideo_63 = wanimagetovideo.EXECUTE_NORMALIZED(width=args.width, height=args.height, length=args.length, batch_size=1, positive=get_value_at_index(cliptextencode_6, 0), negative=get_value_at_index(cliptextencode_7, 0), vae=get_value_at_index(vaeloader_39, 0), start_image=get_value_at_index(loadimage_62, 0))
+        modelsamplingsd3_54 = modelsamplingsd3.patch(shift=8.0, model=get_value_at_index(loraloader_78, 0))
+
+        current_seed = int(time.time() * 1000)
+        random.seed(current_seed)
+        ksampleradvanced_74 = ksampleradvanced.sample(add_noise="enable", noise_seed=random.randint(1, 2**64), steps=4, cfg=1, sampler_name="lcm", scheduler="simple", start_at_step=0, end_at_step=2, return_with_leftover_noise="enable", model=get_value_at_index(modelsamplingsd3_54, 0), positive=get_value_at_index(wanimagetovideo_63, 0), negative=get_value_at_index(wanimagetovideo_63, 1), latent_image=get_value_at_index(wanimagetovideo_63, 2))
+
+        del unetloadergguf_84, loraloader_78, modelsamplingsd3_54, loadimage_62, model, clip
+        clear_memory()
+        print("Step 1 finished and memory cleared.")
+
+        print("Starting Step 2: Refinement Sampling")
+        unetloadergguf_85 = unetloadergguf.load_unet(unet_name="Wan2.2-I2V-A14B-LowNoise-Q4_K_S.gguf")
+
+        # 2번 샘플러 LoRA 체인
+        model, clip = loraloader.load_lora(lora_name="lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank32_bf16.safetensors", strength_model=2, strength_clip=2, model=get_value_at_index(unetloadergguf_85, 0), clip=get_value_at_index(cliploader_38, 0))
+        if args.custom_lora_2_name and args.custom_lora_2_name != "None":
+            print(f"2번 샘플러에 커스텀 LoRA 적용: {args.custom_lora_2_name}")
+            model, clip = loraloader.load_lora(lora_name=args.custom_lora_2_name, strength_model=args.custom_lora_2_strength_model, strength_clip=args.custom_lora_2_strength_clip, model=model, clip=clip)
+        loraloader_86 = loraloader.load_lora(lora_name="FastWan_T2V_14B_480p_lora_rank_128_bf16.safetensors", strength_model=0.5, strength_clip=0.5, model=model, clip=clip)
+
+        modelsamplingsd3_55 = modelsamplingsd3.patch(shift=8, model=get_value_at_index(loraloader_86, 0))
+
+        current_seed_2 = int(time.time() * 1000)
+        random.seed(current_seed_2)
+        ksampleradvanced_75 = ksampleradvanced.sample(add_noise="disable", noise_seed=random.randint(1, 2**64), steps=4, cfg=1, sampler_name="lcm", scheduler="simple", start_at_step=2, end_at_step=10000, return_with_leftover_noise="disable", model=get_value_at_index(modelsamplingsd3_55, 0), positive=get_value_at_index(wanimagetovideo_63, 0), negative=get_value_at_index(wanimagetovideo_63, 1), latent_image=get_value_at_index(ksampleradvanced_74, 0))
+
+        del unetloadergguf_85, loraloader_86, modelsamplingsd3_55, cliploader_38, cliptextencode_6, cliptextencode_7, ksampleradvanced_74, wanimagetovideo_63, model, clip
+        clear_memory()
+        print("Step 2 finished and memory cleared.")
+
+        print("Starting Step 3: VAE Decode and Save")
+        vaedecode_8 = vaedecode.decode(samples=get_value_at_index(ksampleradvanced_75, 0), vae=get_value_at_index(vaeloader_39, 0))
+        saveimage.save_images(filename_prefix="temp/example", images=get_value_at_index(vaedecode_8, 0))
+
+        # ✨ [최적화] 사용이 끝난 VAE 관련 변수 즉시 삭제
+        del ksampleradvanced_75, vaeloader_39, vaedecode_8
+        clear_memory()
+        print("Step 3 finished and memory cleared.")
+
+        if args.upscale_ratio > 1:
+            print("Starting Steps 4 & 5: Upscaling")
+            upscalemodelloader_88 = upscalemodelloader.load_model(model_name="2x-AnimeSharpV4_Fast_RCAN_PU.safetensors")
+
+            # Upscale Part 1
+            vhs_loadimagespath_96 = vhs_loadimagespath.load_images(directory=f"{COMFYUI_BASE_PATH}/output/temp", image_load_cap=40)
+            if get_value_at_index(vhs_loadimagespath_96, 0) is not None and get_value_at_index(vhs_loadimagespath_96, 0).shape[0] > 0:
+                imageupscalewithmodel_92 = imageupscalewithmodel.upscale(upscale_model=get_value_at_index(upscalemodelloader_88, 0), image=get_value_at_index(vhs_loadimagespath_96, 0))
+                mxslider_91 = mxslider.main(Xi=args.upscale_ratio, Xf=args.upscale_ratio, isfloatX=1)
+                easy_mathfloat_90 = easy_mathfloat.float_math_operation(a=get_value_at_index(mxslider_91, 0), b=2, operation="divide")
+                imagescaleby_93 = imagescaleby.upscale(upscale_method="nearest-exact", scale_by=get_value_at_index(easy_mathfloat_90, 0), image=get_value_at_index(imageupscalewithmodel_92, 0))
+                saveimage.save_images(filename_prefix="up/example", images=get_value_at_index(imagescaleby_93, 0))
+                del vhs_loadimagespath_96, imageupscalewithmodel_92, imagescaleby_93, easy_mathfloat_90, mxslider_91
+                clear_memory()
+
+            # Upscale Part 2
+            vhs_loadimagespath_98 = vhs_loadimagespath.load_images(directory=f"{COMFYUI_BASE_PATH}/output/temp", skip_first_images=40)
+            if get_value_at_index(vhs_loadimagespath_98, 0) is not None and get_value_at_index(vhs_loadimagespath_98, 0).shape[0] > 0:
+                imageupscalewithmodel_100 = imageupscalewithmodel.upscale(upscale_model=get_value_at_index(upscalemodelloader_88, 0), image=get_value_at_index(vhs_loadimagespath_98, 0))
+                mxslider_102 = mxslider.main(Xi=args.upscale_ratio, Xf=args.upscale_ratio, isfloatX=1)
+                easy_mathfloat_103 = easy_mathfloat.float_math_operation(a=get_value_at_index(mxslider_102, 0), b=2, operation="divide")
+                imagescaleby_101 = imagescaleby.upscale(upscale_method="nearest-exact", scale_by=get_value_at_index(easy_mathfloat_103, 0), image=get_value_at_index(imageupscalewithmodel_100, 0))
+                saveimage.save_images(filename_prefix="up/example", images=get_value_at_index(imagescaleby_101, 0))
+                del vhs_loadimagespath_98, imageupscalewithmodel_100, easy_mathfloat_103, imagescaleby_101, mxslider_102
+
+            del upscalemodelloader_88
+            clear_memory()
+            print("Upscaling finished and memory cleared.")
+            rife_input_dir = f"{COMFYUI_BASE_PATH}/output/up"
+        else:
+            print("Skipping Upscaling (Upscale ratio <= 1)")
+            rife_input_dir = f"{COMFYUI_BASE_PATH}/output/temp"
+
+        print("Starting Step 6: RIFE and Combine")
+        vhs_loadimagespath_97 = vhs_loadimagespath.load_images(directory=rife_input_dir)
+        rife_vfi_94 = rife_vfi.vfi(ckpt_name="rife47.pth", multiplier=2, frames=get_value_at_index(vhs_loadimagespath_97, 0))
+        vhs_videocombine.combine_video(frame_rate=32, loop_count=0, filename_prefix="AnimateDiff", format="video/h264-mp4", pix_fmt="yuv420p", crf=19, save_metadata=True, trim_to_audio=False, pingpong=False, save_output=True, images=get_value_at_index(rife_vfi_94, 0))
+
+        # ✨ [최적화] 마지막 단계에서 사용한 변수 삭제
+        del vhs_loadimagespath_97, rife_vfi_94
+        clear_memory()
+        print("✅ All steps finished.")
+
+        output_dir = os.path.join(COMFYUI_BASE_PATH, "output")
+        video_files = glob.glob(os.path.join(output_dir, '**', '*.mp4'), recursive=True)
+
+        if not video_files:
+            raise FileNotFoundError("생성된 동영상 파일을 찾을 수 없습니다!")
+
+        latest_video = max(video_files, key=os.path.getctime)
+        print(f"LATEST_VIDEO_PATH:{latest_video}")
+
+if __name__ == "__main__":
+    main()