Update worker_runpod.py

worker_runpod.py

@@ -1,170 +1,109 @@
-import json
-import os
-import runpod
-import numpy as np
+import os, json, requests, random, runpod
+
 import torch
-import requests
-import uuid
-from diffusers import (AutoencoderKL, CogVideoXDDIMScheduler, DDIMScheduler,
-                       DPMSolverMultistepScheduler,
-                       EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
-                       PNDMScheduler)
-from transformers import T5EncoderModel, T5Tokenizer
-from omegaconf import OmegaConf
-from PIL import Image
-from cogvideox.models.transformer3d import CogVideoXTransformer3DModel
-from cogvideox.models.autoencoder_magvit import AutoencoderKLCogVideoX
-from cogvideox.pipeline.pipeline_cogvideox import CogVideoX_Fun_Pipeline
-from cogvideox.pipeline.pipeline_cogvideox_inpaint import CogVideoX_Fun_Pipeline_Inpaint
+from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
 from cogvideox.utils.lora_utils import merge_lora, unmerge_lora
-from cogvideox.utils.utils import get_image_to_video_latent, save_videos_grid
-from cogvideox.data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio
-from huggingface_hub import HfApi, HfFolder
-
-tokenxf = os.getenv("HF_API_TOKEN")
-# Low GPU memory mode
-low_gpu_memory_mode = False
-lora_path = "/content/shirtlift.safetensors"
-
-def to_pil(image):
-    if isinstance(image, Image.Image):
-        return image
-    if isinstance(image, torch.Tensor):
-        return tensor2pil(image)
-    if isinstance(image, np.ndarray):
-        return numpy2pil(image)
-    raise ValueError(f"Cannot convert {type(image)} to PIL.Image")
-
-
-def download_image(url, download_dir="asset"):
-    # Ensure the download directory exists
-    if not os.path.exists(download_dir):
-        os.makedirs(download_dir, exist_ok=True)
-    
-    # Send the request and check for successful response
-    response = requests.get(url, stream=True)
-    if response.status_code == 200:
-        # Determine file extension based on content type
-        content_type = response.headers.get("Content-Type")
-        if content_type == "image/png":
-            ext = "png"
-        elif content_type == "image/jpeg":
-            ext = "jpg"
-        else:
-            ext = "jpg"  # default to .jpg if content type is unrecognized
-        
-        # Generate a random filename with the correct extension
-        filename = f"{uuid.uuid4().hex}.{ext}"
-        file_path = os.path.join(download_dir, filename)
-        
-        # Save the image
-        with open(file_path, "wb") as f:
-            for chunk in response.iter_content(1024):
-                f.write(chunk)
-        
-        print(f"Image downloaded to {file_path}")
-        return file_path
-    else:
-        raise Exception(f"Failed to download image from {url}, status code: {response.status_code}")
-
-# Usage
-# validation_image_start = values.get("validation_image_start", "https://example.com/path/to/image.png")
-# downloaded_image_path = download_image(validation_image_start)
-model_id = "/content/model"
-transformer = CogVideoXTransformer3DModel.from_pretrained_2d(model_id, subfolder="transformer").to(torch.bfloat16)
-
-vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae").to(torch.bfloat16)
-
-text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder")
-
-sampler_dict = {
-    "Euler": EulerDiscreteScheduler,
-    "Euler A": EulerAncestralDiscreteScheduler,
-    "DPM++": DPMSolverMultistepScheduler,
-    "PNDM": PNDMScheduler,
-    "DDIM_Cog": CogVideoXDDIMScheduler,
-    "DDIM_Origin": DDIMScheduler,
-}
-scheduler = sampler_dict["DPM++"].from_pretrained(model_id, subfolder="scheduler")
-
-# Pipeline setup
-if transformer.config.in_channels != vae.config.latent_channels:
-    pipeline = CogVideoX_Fun_Pipeline_Inpaint.from_pretrained(
-        model_id, vae=vae, text_encoder=text_encoder,
-        transformer=transformer, scheduler=scheduler,
-        torch_dtype=torch.bfloat16
-    )
-else:
-    pipeline = CogVideoX_Fun_Pipeline.from_pretrained(
-        model_id, vae=vae, text_encoder=text_encoder,
-        transformer=transformer, scheduler=scheduler,
-        torch_dtype=torch.bfloat16
-    )
+from diffusers.utils import export_to_video, load_image
+from transformers import T5EncoderModel, T5Tokenizer
 
-if low_gpu_memory_mode:
-    pipeline.enable_sequential_cpu_offload()
-else:
-    pipeline.enable_model_cpu_offload()
+with torch.inference_mode():
+    model_id = "/content/model"
+    transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16)
+    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16)
+    vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16)
+    tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
+    pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16).to("cuda")
+    # pipe.enable_model_cpu_offload()
+
+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"]
-    prompt = values["prompt"]
-    negative_prompt = values.get("negative_prompt", "blurry, blurred, blurry face")
-    guidance_scale = values.get("guidance_scale", 6.0)
-    seed = values.get("seed", 42)
-    num_inference_steps = values.get("num_inference_steps", 18)
-    base_resolution = values.get("base_resolution", 512)
-    
-    video_length = values.get("video_length", 53)
-    fps = values.get("fps", 10)
-    lora_weight = values.get("lora_weight", 1.00)
-    save_path = "samples"
-    partial_video_length = values.get("partial_video_length", None)
-    overlap_video_length = values.get("overlap_video_length", 4)
-    validation_image_start = values.get("validation_image_start", "asset/1.png")
-    downloaded_image_path = download_image(validation_image_start)
-    validation_image_end = values.get("validation_image_end", None)
-    
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    if lora_path is not None:
-        pipeline = merge_lora(pipeline, lora_path, lora_weight)
-    
-    aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
-    start_img = Image.open(downloaded_image_path)
-    original_width, original_height = start_img.size
-    closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size)
-    height, width = [int(x / 16) * 16 for x in closest_size]
-    sample_size = [height, width]
-    video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
-    input_video, input_video_mask, clip_image = get_image_to_video_latent(downloaded_image_path, validation_image_end, video_length=video_length, sample_size=sample_size)
-        
-    with torch.no_grad():
-        sample = pipeline(prompt=prompt,num_frames=video_length,negative_prompt=negative_prompt,height=sample_size[0],width=sample_size[1],generator=generator,guidance_scale=guidance_scale,num_inference_steps=num_inference_steps,video=input_video,mask_video=input_video_mask).videos        
-    
-    if not os.path.exists(save_path):
-        os.makedirs(save_path, exist_ok=True)
-
-    index = len([path for path in os.listdir(save_path)]) + 1
-    prefix = str(index).zfill(8)
-    video_path = os.path.join(save_path, f"{prefix}.mp4")
-    save_videos_grid(sample, video_path, fps=fps)
-
-
-    hf_api = HfApi()
-    repo_id = "meepmoo/h4h4jejdf"  # Set your HF repo
-    hf_api.upload_file(
-       path_or_fileobj=video_path,
-       path_in_repo=f"{prefix}.mp4",
-       repo_id=repo_id,
-       token=tokenxf,
-       repo_type="model"
-    )
-
-    
-    result_url = f"https://huggingface.co/{repo_id}/blob/main/{prefix}.mp4"
-    result_url = ""
-    job_id = values.get("job_id", "default-job-id")  # For RunPod job tracking
-    return {"jobId": job_id, "result": result_url, "status": "DONE"}
+    lora_path = "/content/shirtlift.safetensors"
+    lora_weight = 1.0
+    pipe = merge_lora(pipe, lora_path, lora_weight)
+    input_image = values['input_image_check']
+    input_image = download_file(url=input_image, save_dir='/content/input', file_name='input_image_tost')
+    prompt = values['prompt']
+    # guidance_scale = values['guidance_scale']
+    # use_dynamic_cfg = values['use_dynamic_cfg']
+    # num_inference_steps = values['num_inference_steps']
+    # fps = values['fps']
+    guidance_scale = 6
+    use_dynamic_cfg = True
+    num_inference_steps = 50
+    fps = 8
+
+    image = load_image(input_image)
+    video = pipe(image=image, prompt=prompt, guidance_scale=guidance_scale, use_dynamic_cfg=use_dynamic_cfg, num_inference_steps=num_inference_steps).frames[0]
+    export_to_video(video, "/content/cogvideox_5b_i2v_tost.mp4", fps=fps)
+
+    result = "/content/cogvideox_5b_i2v_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})