pool

baseline_utils.py

@@ -123,38 +123,40 @@ def scenes_caption(scenes, api_key):
 
     return "\n\n".join(captions)
 
+# Define the single video generation function in the global scope
+def generate_single_video(gpu_id, prompt, writer_description, fps, i):
+    # Assign the specific GPU for this process
+    device = f"cuda:{gpu_id}"
+
+    # Initialize the pipeline for this GPU
+    pipe = CogVideoXPipeline.from_pretrained(
+        "THUDM/CogVideoX-5b",
+        torch_dtype=torch.bfloat16,
+        cache_dir="./CogVideoX-5b"
+    )
+    pipe.to(device)  # Move the model to the assigned GPU
+    pipe.enable_model_cpu_offload()
+    pipe.vae.enable_tiling()
 
-@spaces.GPU
-def generate_video(scene_list, writer_description, opt, fps=24):  # Lower fps
+    # Generate the video
+    video = pipe(
+        prompt=prompt + f'\nThe main character is described as: {writer_description}.',
+        num_videos_per_prompt=1,
+        num_inference_steps=50,
+        num_frames=fps,
+        guidance_scale=6,
+        generator=torch.Generator(device=device).manual_seed(42),
+    ).frames[0]
 
-    # Function to generate a single video
-    def generate_single_video(gpu_id, prompt, writer_description, fps, i):
-        # Assign the specific GPU for this process
-        device = f"cuda:{gpu_id}"
+    # Save the video
+    video_path = export_to_video(video, output_video_path=f'videos/video{i}.mp4')
+    return video_path
 
-        # Initialize the pipeline for this GPU
-        pipe = CogVideoXPipeline.from_pretrained(
-            "THUDM/CogVideoX-5b",
-            torch_dtype=torch.bfloat16,
-            cache_dir="./CogVideoX-5b"
-        )
-        pipe.to(device)  # Move the model to the assigned GPU
-        pipe.enable_model_cpu_offload()
-        pipe.vae.enable_tiling()
-
-        # Generate the video
-        video = pipe(
-            prompt=prompt + f'\nThe main character is described as: {writer_description}.',
-            num_videos_per_prompt=1,
-            num_inference_steps=50,
-            num_frames=fps,
-            guidance_scale=6,
-            generator=torch.Generator(device=device).manual_seed(42),
-        ).frames[0]
-
-        # Save the video
-        video_path = export_to_video(video, output_video_path=f'videos/video{i}.mp4')
-        return video_path
+
+@spaces.GPU
+def generate_video(scene_list, writer_description, opt, fps=24):  # Lower fps
+    # Set TOKENIZERS_PARALLELISM to avoid tokenizer warnings
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
     # Ensure the output directory exists
     os.makedirs("videos", exist_ok=True)
@@ -199,6 +201,7 @@ def concatenate_videos_music(video_paths, output_path, audio_path):
     # Write the concatenated video to a file
     final_clip.write_videofile(output_path, codec="libx264", audio_codec="aac")
 
+@spaces.GPU
 def summarizer_for_audio(input_text):
   if torch.cuda.is_available():
       device = "cuda"