Spaces:
Runtime error
Runtime error
File size: 3,854 Bytes
0328207 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from tqdm import tqdm
from torchvision.io import write_video
from torch.utils.data import DataLoader, SequentialSampler
from stream_inference_wrapper import StreamInferenceWrapper
# from stream_drag_inference_wrapper import StreamDragInferenceWrapper
from utils.dataset import TextDataset
from utils.misc import set_seed
from hydra import initialize, compose
from hydra.core.global_hydra import GlobalHydra
from demo_utils.memory import gpu, get_cuda_free_memory_gb
def main():
output_block_number = 27
config_dir = "configs"
stream_config_name = "self_forcing_dmd_vsink_stream"
# stream_config_name = "self_forcing_dmd_vsink_stream_drag"
data_path = "prompts/MovieGenVideoBench_extended.txt"
seed = 42
set_seed(seed)
output_folder = "outputs-stream"
output_folder = f"{output_folder}/blk{output_block_number}-{stream_config_name}-seed{seed}"
print(f"Free VRAM {get_cuda_free_memory_gb(gpu)} GB")
# low_memory = get_cuda_free_memory_gb(gpu) < 40
# Create dataset
dataset = TextDataset(prompt_path=data_path)
num_prompts = len(dataset)
print(f"Number of prompts: {num_prompts}")
sampler = SequentialSampler(dataset)
dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
os.makedirs(output_folder, exist_ok=True)
if GlobalHydra.instance().is_initialized():
GlobalHydra.instance().clear()
with initialize(version_base=None, config_path=config_dir):
stream_config = compose(config_name=stream_config_name)
print(f"{stream_config = }")
stream_inference = StreamInferenceWrapper(
stream_model_config=stream_config,
checkpoint_path="./checkpoints/self_forcing_dmd.pt",
total_generate_block_number=output_block_number,
use_ema=True,
seed=seed,
)
for i, batch_data in tqdm(enumerate(dataloader)):
idx = batch_data["idx"].item()
print(f"{idx = }")
# For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
# Unpack the batch data for convenience
if isinstance(batch_data, dict):
batch = batch_data
elif isinstance(batch_data, list):
batch = batch_data[0] # First (and only) item in the batch
# For text-to-video, batch is just the text prompt
prompt = batch["prompts"][0]
print(f"{prompt = }")
extended_prompt = batch["extended_prompts"][0] if "extended_prompts" in batch else None
print(f"{extended_prompt = }")
set_seed(seed)
stream_inference.reset()
current_block_index = 0
block_step = 3
while current_block_index < output_block_number:
end_block_index = current_block_index + block_step
all_video, current_video = stream_inference.inference(
start_block_index=current_block_index,
end_block_index=end_block_index,
prompt=prompt,
)
# Save the video if the current prompt is not a dummy prompt
if idx < num_prompts:
current_video_output_path = os.path.join(
output_folder,
f"{idx:04d}-{prompt[:50].replace(' ', '_')}-{current_block_index:02d}-{end_block_index:02d}.mp4",
)
write_video(current_video_output_path, current_video, fps=16)
all_video_output_path = os.path.join(
output_folder,
f"{idx:04d}-{prompt[:50].replace(' ', '_')}-{0:02d}-{end_block_index:02d}.mp4",
)
write_video(all_video_output_path, all_video, fps=16)
current_block_index = end_block_index
if __name__ == "__main__":
main()
|