Update video_generator.py

video_generator.py

@@ -1,15 +1,13 @@
 import os
 import numpy as np
+import logging
 from PIL import Image
 from moviepy.editor import ImageSequenceClip, AudioFileClip
 from moviepy.config import change_settings
-import logging
+import torch
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
-# Configure MoviePy
 change_settings({"IMAGEMAGICK_BINARY": "/usr/bin/convert"})
 
 def generate_video_pipeline(
@@ -22,54 +20,61 @@ def generate_video_pipeline(
     fps: int = 24
 ):
     os.makedirs(output_dir, exist_ok=True)
-    image_path = None  # Initialize variable
+    image_path = None
     audio_path = None
 
     try:
-        # Generate initial image
-        logger.info(f"Generating image for prompt: {prompt}")
+        # Generate image with device awareness
         image_result = text_to_image_model(prompt)
         image = image_result.images[0]
+        
+        # Convert to CPU-compatible format if needed
+        if image_to_video_model.device.type == "cpu":
+            image = image.float()
+            
         image_path = os.path.join(output_dir, "frame.png")
         image.save(image_path)
 
-        # Generate video frames
-        logger.info(f"Generating {duration}s video @ {fps}fps")
+        # Video generation
         video_frames = image_to_video_model(
             prompt,
             image=image,
             num_frames=int(duration * fps),
-            num_inference_steps=25
+            num_inference_steps=25,
+            generator=torch.Generator(
+                device=image_to_video_model.device
+            )
         ).frames
 
-        # Create video clip
+        # Create video
         video_path = os.path.join(output_dir, "output.mp4")
-        clip = ImageSequenceClip([np.array(frame) for frame in video_frames], fps=fps)
+        clip = ImageSequenceClip(
+            [np.array(frame) for frame in video_frames], 
+            fps=fps
+        )
 
-        # Add audio if available
+        # Audio handling
         if tts_model and prompt:
-            logger.info("Generating audio track")
             audio = tts_model.generate(prompt)
             audio_path = os.path.join(output_dir, "audio.wav")
             audio.save(audio_path)
-            audio_clip = AudioFileClip(audio_path)
-            clip = clip.set_audio(audio_clip)
+            clip = clip.set_audio(AudioFileClip(audio_path))
 
-        # Export final video
-        clip.write_videofile(video_path, codec="libx264", audio_codec="aac", logger=None)
-        logger.info(f"Video generated at {video_path}")
+        clip.write_videofile(video_path, 
+                           codec="libx264", 
+                           audio_codec="aac",
+                           logger=None)
         return video_path
 
     except Exception as e:
-        logger.error(f"Pipeline error: {str(e)}", exc_info=True)
-        raise RuntimeError(f"Video generation failed: {str(e)}") from e
+        logger.error(f"Generation failed: {str(e)}")
+        raise
 
     finally:
-        # Cleanup temporary files safely
+        # Safe cleanup
         for path in [image_path, audio_path]:
             if path and os.path.exists(path):
                 try:
                     os.remove(path)
-                    logger.debug(f"Cleaned up temporary file: {path}")
-                except Exception as cleanup_error:
-                    logger.warning(f"Failed to clean {path}: {str(cleanup_error)}")
+                except Exception as e:
+                    logger.warning(f"Cleanup failed: {str(e)}")