Update worker_runpod.py

worker_runpod.py

@@ -3,6 +3,8 @@ import os
 import runpod
 import numpy as np
 import torch
+import requests
+import uuid
 from diffusers import (AutoencoderKL, CogVideoXDDIMScheduler, DDIMScheduler,
                        DPMSolverMultistepScheduler,
                        EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
@@ -22,6 +24,29 @@ tokenxf = os.getenv("HF_API_TOKEN")
 # Low GPU memory mode
 low_gpu_memory_mode = False
 
+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)
+    
+    # Generate a random filename with a .png extension
+    filename = f"{uuid.uuid4().hex}.png"
+    file_path = os.path.join(download_dir, filename)
+    
+    # Download the image and save it
+    response = requests.get(url, stream=True)
+    if response.status_code == 200:
+        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 loading section
 model_id = "/content/model"
 transformer = CogVideoXTransformer3DModel.from_pretrained_2d(
@@ -82,11 +107,12 @@ def generate(input):
     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)
     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(validation_image_start)
+    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]
@@ -97,7 +123,7 @@ def generate(input):
     else:
         # Standard video generation
         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(validation_image_start, validation_image_end, video_length=video_length, sample_size=sample_size)
+        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(