Batch size

[ahmeddawy]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @ahmeddawy, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the training and evaluation infrastructure by introducing robust support for batch processing across the diffusion model pipeline. It refactors data loading and preprocessing to natively handle torch.Tensor inputs, adds configurable batch size arguments, and integrates a dedicated validation loop into the training process. These changes aim to improve training efficiency, provide better tools for monitoring model performance, and make the system more adaptable to various data input formats.

Highlights

• Batch Processing Support: Core data preprocessing functions (preprocess_image, preprocess_video) now directly accept and correctly handle torch.Tensor inputs, enabling efficient batch processing without requiring conversion from PIL Images.
• Validation Loop Integration: A comprehensive validation mechanism has been added to the training runner, allowing for periodic evaluation of the model during training with configurable batch sizes and datasets, improving model monitoring capabilities.
• Configurable Batch Sizes: New command-line arguments (--batch_size, --val_batch_size) have been introduced to control batch sizes for both training and validation, providing greater flexibility in resource utilization.
• Batch-Aware Data Handling: Custom collate functions and input parsing logic (_collate_batch, pick_first_frame, pick_last_frame) ensure that various data types (images, videos, masks) are correctly batched and prepared for model input, streamlining the data pipeline.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for batched training and validation, which is a significant enhancement. The changes are extensive, touching data preprocessing, argument parsing, the training runner, and model implementations to handle batched tensors. The addition of a validation loop with periodic evaluation is a great feature for monitoring model performance during training.

I've found a critical issue in the new batch collation logic that would prevent batched training from working correctly with custom dataset objects. I've also included a suggestion to improve the efficiency and readability of a new test script.

Overall, this is a great contribution. Addressing the identified issues will make the new batching functionality robust.

[gemini-code-assist]

The check isinstance(value, list) is too restrictive. It will likely fail for custom sequence-like objects such as VideoData returned by the dataset, causing an error during batch collation. To make this more robust, you should check for the general Sequence type from collections.abc and explicitly exclude strings.

You'll need to add from collections.abc import Sequence at the top of the file.

Suggested change

	if isinstance(value, list):
	if isinstance(value, Sequence) and not isinstance(value, str):

[gemini-code-assist]

Similar to the previous comment, this isinstance(value, list) check is too restrictive and will fail for custom sequence types like VideoData. Using collections.abc.Sequence will make the collate function more general and prevent errors with different dataset implementations.

Suggested change

	if isinstance(value, list):
	if isinstance(value, Sequence) and not isinstance(value, str):

[gemini-code-assist]

This block of code for loading and preparing the video, reference image, and mask can be made more efficient and readable.

1. The if/else block to truncate the video is verbose. Slicing ([:81]) is more Pythonic.
2. reference_image is created by reloading the same video file, which is inefficient. It can be retrieved from the VideoData object that's already loaded.
3. The list comprehension to create the mask can also be simplified using slicing.

The suggested change refactors this logic to be more concise and avoid redundant file I/O.

Suggested change

	video = VideoData(str(video_path), height=480, width=832)
	if len(video) >= 81 :
	video = [video[i] for i in range(81)]
	else:
	video = [video[i] for i in range(len(video))]
	reference_image = VideoData(str(video_path), height=480, width=832)[0]
	vace_mask = VideoData(str(mask_path), height=480, width=832)
	vace_mask = [vace_mask[i] for i in range(len(video))]
	video_data = VideoData(str(video_path), height=480, width=832)
	video = list(video_data[:81])
	reference_image = video_data[0]
	vace_mask_data = VideoData(str(mask_path), height=480, width=832)
	vace_mask = list(vace_mask_data[:len(video)])