Update worker_runpod.py

worker_runpod.py

@@ -71,22 +71,22 @@ def download_image(url, download_dir="/content"):
 # 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 = "/runpod-volume/model"
-transformer = CogVideoXTransformer3DModel.from_pretrained_2d(model_id, subfolder="transformer").to(torch.bfloat16)
+with torch.inference_mode():
+    model_id = "/runpod-volume/model"
+    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 = {
+    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")
+    }
+    scheduler = sampler_dict["DPM++"].from_pretrained(model_id, subfolder="scheduler")
 
 # Pipeline setup
 if transformer.config.in_channels != vae.config.latent_channels:
@@ -102,10 +102,12 @@ else:
         torch_dtype=torch.bfloat16
     )
 
-if low_gpu_memory_mode:
-    pipeline.enable_sequential_cpu_offload()
-else:
-    pipeline.enable_model_cpu_offload()
+# if low_gpu_memory_mode:
+#     pipeline.enable_sequential_cpu_offload()
+# else:
+#     pipeline.enable_model_cpu_offload()
+
+
 
 @torch.inference_mode()
 def generate(input):
@@ -118,10 +120,11 @@ 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)
-    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)
@@ -132,8 +135,7 @@ def generate(input):
     
     generator = torch.Generator(device="cuda").manual_seed(seed)
     print("Generator started")
-    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)