Create worker_runpod.py

worker_runpod.py

@@ -0,0 +1,138 @@
+import os, json, requests, random, time, runpod
+from urllib.parse import urlsplit
+
+import torch
+from PIL import Image
+import numpy as np
+
+import asyncio
+import execution
+import server
+loop = asyncio.new_event_loop()
+asyncio.set_event_loop(loop)
+server_instance = server.PromptServer(loop)
+execution.PromptQueue(server)
+
+from nodes import load_custom_node
+from nodes import NODE_CLASS_MAPPINGS
+
+load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-CogVideoXWrapper")
+load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-VideoHelperSuite")
+load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-KJNodes")
+
+LoadImage = NODE_CLASS_MAPPINGS["LoadImage"]()
+ImageResizeKJ = NODE_CLASS_MAPPINGS["ImageResizeKJ"]()
+CogVideoImageEncode = NODE_CLASS_MAPPINGS["CogVideoImageEncode"]()
+CogVideoLoraSelect = NODE_CLASS_MAPPINGS["CogVideoLoraSelect"]()
+DownloadAndLoadCogVideoModel = NODE_CLASS_MAPPINGS["DownloadAndLoadCogVideoModel"]()
+CogVideoTextEncode = NODE_CLASS_MAPPINGS["CogVideoTextEncode"]()
+CLIPLoader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
+CogVideoSampler = NODE_CLASS_MAPPINGS["CogVideoSampler"]()
+CogVideoDecode = NODE_CLASS_MAPPINGS["CogVideoDecode"]()
+VHS_VideoCombine = NODE_CLASS_MAPPINGS["VHS_VideoCombine"]()
+
+with torch.inference_mode():
+    lora = CogVideoLoraSelect.getlorapath("orbit_up_lora_weights.safetensors", 1.0, fuse_lora=True)[0]
+    pipeline = DownloadAndLoadCogVideoModel.loadmodel("THUDM/CogVideoX-5b-I2V", "bf16", fp8_transformer="disabled", compile="disabled", enable_sequential_cpu_offload=False, lora=lora)[0]
+    clip = CLIPLoader.load_clip("t5xxl_fp16.safetensors", type="sd3")[0]
+
+def download_file(url, save_dir, file_name):
+    os.makedirs(save_dir, exist_ok=True)
+    original_file_name = url.split('/')[-1]
+    _, original_file_extension = os.path.splitext(original_file_name)
+    file_path = os.path.join(save_dir, file_name + original_file_extension)
+    response = requests.get(url)
+    response.raise_for_status()
+    with open(file_path, 'wb') as file:
+        file.write(response.content)
+    return file_path
+
+@torch.inference_mode()
+def generate(input):
+    values = input["input"]
+
+    input_image=values['input_image_check']
+    input_image=download_file(url=input_image, save_dir='/content/ComfyUI/input', file_name='input_image')
+    prompt = values['prompt']
+    negative_prompt = values['negative_prompt']
+    seed = values['seed']
+    steps = values['steps']
+    cfg = values['cfg']
+
+    if seed == 0:
+        random.seed(int(time.time()))
+        seed = random.randint(0, 18446744073709551615)
+
+    positive = CogVideoTextEncode.process(clip, prompt, strength=1.0, force_offload=True)[0]
+    negative = CogVideoTextEncode.process(clip, negative_prompt, strength=1.0, force_offload=True)[0]
+
+    image, _ = LoadImage.load_image(input_image)
+    image = ImageResizeKJ.resize(image, width=720, height=480, keep_proportion=False, upscale_method="lanczos", divisible_by=16, crop="center")[0]
+    image_cond_latents = CogVideoImageEncode.encode(pipeline, image, chunk_size=16, enable_tiling=True)[0]
+    samples = CogVideoSampler.process(pipeline, positive, negative, steps, cfg, seed, height=480, width=720, num_frames=49, scheduler="CogVideoXDPMScheduler", denoise_strength=1.0, image_cond_latents=image_cond_latents)
+    frames = CogVideoDecode.decode(samples[0], samples[1], enable_vae_tiling=True, tile_sample_min_height=240, tile_sample_min_width=360, tile_overlap_factor_height=0.2, tile_overlap_factor_width=0.2, auto_tile_size=True)[0]
+
+    out_video = VHS_VideoCombine.combine_video(images=frames, frame_rate=8, loop_count=0, filename_prefix="CogVideoX-I2V", format="video/h264-mp4", save_output=True)
+    source = out_video["result"][0][1][1]
+    destination = f"/content/ComfyUI/output/cogvideox-5b-i2v-dimensionx-{seed}-tost.mp4"
+    shutil.move(source, destination)
+
+    result = f"/content/ComfyUI/output/cogvideox-5b-i2v-dimensionx-{seed}-tost.mp4"
+    try:
+        notify_uri = values['notify_uri']
+        del values['notify_uri']
+        notify_token = values['notify_token']
+        del values['notify_token']
+        discord_id = values['discord_id']
+        del values['discord_id']
+        if(discord_id == "discord_id"):
+            discord_id = os.getenv('com_camenduru_discord_id')
+        discord_channel = values['discord_channel']
+        del values['discord_channel']
+        if(discord_channel == "discord_channel"):
+            discord_channel = os.getenv('com_camenduru_discord_channel')
+        discord_token = values['discord_token']
+        del values['discord_token']
+        if(discord_token == "discord_token"):
+            discord_token = os.getenv('com_camenduru_discord_token')
+        job_id = values['job_id']
+        del values['job_id']
+        default_filename = os.path.basename(result)
+        with open(result, "rb") as file:
+            files = {default_filename: file.read()}
+        payload = {"content": f"{json.dumps(values)} <@{discord_id}>"}
+        response = requests.post(
+            f"https://discord.com/api/v9/channels/{discord_channel}/messages",
+            data=payload,
+            headers={"Authorization": f"Bot {discord_token}"},
+            files=files
+        )
+        response.raise_for_status()
+        result_url = response.json()['attachments'][0]['url']
+        notify_payload = {"jobId": job_id, "result": result_url, "status": "DONE"}
+        web_notify_uri = os.getenv('com_camenduru_web_notify_uri')
+        web_notify_token = os.getenv('com_camenduru_web_notify_token')
+        if(notify_uri == "notify_uri"):
+            requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+        else:
+            requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+            requests.post(notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token})
+        return {"jobId": job_id, "result": result_url, "status": "DONE"}
+    except Exception as e:
+        error_payload = {"jobId": job_id, "status": "FAILED"}
+        try:
+            if(notify_uri == "notify_uri"):
+                requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+            else:
+                requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+                requests.post(notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token})
+        except:
+            pass
+        return {"jobId": job_id, "result": f"FAILED: {str(e)}", "status": "FAILED"}
+    finally:
+        if os.path.exists(result):
+            os.remove(result)
+        if os.path.exists(input_image):
+            os.remove(input_image)
+
+runpod.serverless.start({"handler": generate})