kenken999 commited on
Commit
d510900
1 Parent(s): 8f00609
Files changed (1) hide show
  1. routers/gra_06_video/video.py +15 -0
routers/gra_06_video/video.py CHANGED
@@ -31,11 +31,26 @@ pipe.enable_vae_slicing()
31
 
32
 
33
  def to_video(frames: list[np.ndarray], fps: int) -> str:
 
34
  out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
 
 
35
  writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps)
 
 
36
  for frame in frames:
 
 
 
 
 
 
 
37
  writer.append_data(frame)
 
 
38
  writer.close()
 
39
  return out_file.name
40
 
41
 
 
31
 
32
 
33
  def to_video(frames: list[np.ndarray], fps: int) -> str:
34
+ # Ensure output file is created and will be kept after closing
35
  out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
36
+
37
+ # Initialize video writer
38
  writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps)
39
+
40
+ # Process each frame
41
  for frame in frames:
42
+ # Ensure frame has 1, 2, 3, or 4 channels
43
+ if len(frame.shape) == 2: # If the frame is grayscale
44
+ frame = np.stack([frame] * 3, axis=-1) # Convert to RGB
45
+ elif frame.shape[2] not in [1, 2, 3, 4]:
46
+ raise ValueError(f"Frame must have 1, 2, 3, or 4 channels, not {frame.shape[2]}")
47
+
48
+ # Append frame to video
49
  writer.append_data(frame)
50
+
51
+ # Close the writer and release the file
52
  writer.close()
53
+
54
  return out_file.name
55
 
56