worker_runpod.py

import json
import os
import runpod
import numpy as np
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 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"
weight_dtype = torch.bfloat16
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="/content"):
    # 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}")
        os.chmod(file_path, 0o777)
        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)
with torch.inference_mode():
    model_id = "/runpod-volume/model"
    transformer = CogVideoXTransformer3DModel.from_pretrained_2d(
        model_id, subfolder="transformer"
    ).to(weight_dtype)

    vae = AutoencoderKLCogVideoX.from_pretrained(
        model_id, subfolder="vae"
    ).to(weight_dtype)

    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=weight_dtype)

    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")


    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=weight_dtype
        )
    else:
        pipeline = CogVideoX_Fun_Pipeline.from_pretrained(
            model_id, 
            vae=vae, 
            text_encoder=text_encoder,
            transformer=transformer, 
            scheduler=scheduler,
            torch_dtype=weight_dtype
        )

    pipeline = merge_lora(pipeline, lora_path, 1.00)


if low_gpu_memory_mode:
    pipeline.enable_sequential_cpu_offload()
else:
    pipeline.enable_model_cpu_offload()



@torch.inference_mode()
def generate(input):
    values = input["input"]
    prompt = values["prompt"]
    print("starting Generate function")
    print(prompt)
    negative_prompt = values.get("negative_prompt", "The video is not of a high quality, it has a low resolution. Watermark present in each frame. Strange motion trajectory. blurry, blurred, grainy, distortion, 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", 49)
    fps = values.get("fps", 10)
    
    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")
    print(validation_image_start)
    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)
    print("Generator started")

    
    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]
    print("Getting closest ratio")
    print(closest_ratio)
    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)
    filename2 = f"{uuid.uuid4().hex}.mp4"
    video_path = os.path.join(save_path, filename2)
    save_videos_grid(sample, video_path, fps=fps)

    print("Video saved to grid, uploading to huggingface")
    hf_api = HfApi()
    
    repo_id = "meepmoo/h4h4jejdf"  # Set your HF repo
    hf_api.upload_file(
       path_or_fileobj=video_path,
       path_in_repo=filename2,
       repo_id=repo_id,
       token=tokenxf,
       repo_type="model"
    )
    
    print("Video uploaded to huggingface returing output")
    result_url = f"https://huggingface.co/{repo_id}/resolve/main/{filename2}"
    job_id = values.get("job_id", "default-job-id")  # For RunPod job tracking
    return {"jobId": job_id, "result": result_url, "status": "DONE"}

runpod.serverless.start({"handler": generate})