Update worker_runpod.py

worker_runpod.py

@@ -73,34 +73,48 @@ def download_image(url, download_dir="/content"):
 # 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(torch.bfloat16)
+    transformer = CogVideoXTransformer3DModel.from_pretrained_2d(
+        model_id, subfolder="transformer"
+    ).to(torch.bfloat16)
 
-    vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae").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,
+        "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")
+    lora_weight = values.get("lora_weight", 1.00)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight)
+
+    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
+        )
 
-# 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
-    )
 
 # if low_gpu_memory_mode:
 #     pipeline.enable_sequential_cpu_offload()
@@ -120,8 +134,7 @@ def generate(input):
     seed = values.get("seed", 42)
     num_inference_steps = values.get("num_inference_steps", 18)
     base_resolution = values.get("base_resolution", 512)
-    lora_weight = values.get("lora_weight", 1.00)
-    pipeline = merge_lora(pipeline, lora_path, lora_weight)
+    
     video_length = values.get("video_length", 53)
     fps = values.get("fps", 10)