pool

baseline_utils.py

@@ -9,10 +9,9 @@ from diffusers.utils import export_to_video
 import os
 import spaces
 from moviepy.editor import VideoFileClip, AudioFileClip, concatenate_videoclips
+from transformers import pipeline
 import requests
 from transformers import pipeline
-from multiprocessing import Pool
-
 
 # Utilize the Google Cloud Vision API to recognize text in the
 # input input_images (diary input_images), https://cloud.google.com/vision.
@@ -123,62 +122,52 @@ 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
+@spaces.GPU
+def generate_video(scene_list, writer_description, opt, fps=24):  # Lower fps
+
     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"
+    # Check for available device: CUDA, MPS, or CPU
+    if torch.cuda.is_available():
+        device = "cuda"
+        print("Using CUDA backend.")
+    elif torch.backends.mps.is_available():
+        device = "mps"
+        print("Using MPS backend.")
+    else:
+        device = "cpu"
+        print("CUDA and MPS not available. Falling back to CPU.")
 
-    # Ensure the output directory exists
+    # Truncate the prompt to fit the CLIP token limit
     os.makedirs("videos", exist_ok=True)
-
-    # Assign each task to a GPU in a round-robin fashion
-    num_gpus = torch.cuda.device_count()
-    if num_gpus < 4:
-        raise RuntimeError("This code assumes at least 4 GPUs are available.")
-
-    tasks = [(i % num_gpus, prompt, writer_description, fps, i) for i, prompt in enumerate(scene_list)]
-
-    # Parallelize using multiprocessing
-    with Pool(processes=num_gpus) as pool:
-        video_paths = pool.starmap(generate_single_video, tasks)
+    video_paths = []
+    for i, prompt in enumerate(scene_list):
+        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]
+
+        video_path = export_to_video(video, output_video_path=f'videos/video{i}.mp4')
+        video_paths.append(video_path)
 
     # Concatenate the generated videos into a single video
     concatenated_video_path = "videos/combined_video_music.mp4"
     if opt == "Narration":
-        concatenate_videos_music(video_paths, concatenated_video_path, audio_path="narration.mp3")
+      concatenate_videos_music(video_paths, concatenated_video_path, audio_path="narration.mp3")
     else:
-        concatenate_videos_music(video_paths, concatenated_video_path, audio_path="meow-meow-meow-tiktok.mp3")
-
+      concatenate_videos_music(video_paths, concatenated_video_path, audio_path="meow-meow-meow-tiktok.mp3")
     return concatenated_video_path