Update worker_runpod.py

worker_runpod.py

@@ -29,16 +29,27 @@ def download_image(url, download_dir="asset"):
     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
+    # 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:
@@ -47,7 +58,6 @@ def download_image(url, download_dir="asset"):
 # 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(
     model_id, subfolder="transformer", torch_dtype=torch.bfloat16