[OmniGen2] add fine-tuning script

[hadipash]

What does this PR do?

Fixes # (issue)

• mindone/data/loader.py: Added a warning specifying that max_rowsize is deprecated in MindSpore 2.6 or above.
• mindone/data/loader.py: Fixed a bug that caused column selection (dataloader.project(project_columns)) to be applied before batch transforms.
• TrainOneStepWrapper now delegates setting training mode (set_train()) to the model rather than setting the entire network to train mode.

Adds # (feature)

• Adds fine-tuning (full and LoRA) script for OmniGen2.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you make sure to update the documentation with your changes? E.g. record bug fixes or new features in What's New. Here are the
  documentation guidelines
• Did you build and run the code without any errors?
• Did you report the running environment (NPU type/MS version) and performance in the doc? (better record it for data loading, model inference, or training tasks)
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[gemini-code-assist]

Summary of Changes

Hello @hadipash, 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 expands the capabilities of OmniGen2 by introducing robust fine-tuning functionalities. It provides the necessary infrastructure, including data loading, model configuration, and training pipelines, to adapt OmniGen2 models to specific tasks using both full-parameter and LoRA approaches. The changes also include crucial MindSpore compatibility fixes and enhancements for better training stability and observability.

Highlights

• OmniGen2 Fine-Tuning Scripts: Introduced new training scripts and configurations for fine-tuning OmniGen2 models, supporting both full-parameter and LoRA (Low-Rank Adaptation) methods. This includes dedicated YAML configuration files and shell scripts for launching training jobs.
• Enhanced Data Handling and Processing: Added new dataset classes (OmniGen2TrainDataset, OmniGen2Collator) and related utilities to efficiently load, preprocess, and tokenize data for fine-tuning, including support for various task types (text-to-image, image editing, in-context generation) and chat templates.
• MindSpore Compatibility and Optimizations: Implemented MindSpore-specific adjustments for complex number operations in rotary embeddings, updated gradient checkpointing logic, and refined the TrainOneStepWrapper to correctly delegate training mode settings to the model. Also, max_rowsize deprecation in MindSpore 2.6+ is now handled with a warning.
• Improved Training Monitoring and Logging: Integrated new callbacks (BinLossCallback, VisualizationCallback) for detailed loss analysis across time bins and visual output of training progress, along with comprehensive logging utilities for better experiment tracking.
• Documentation and Dependency Updates: The README.md has been updated with detailed instructions for setting up and launching fine-tuning, including configuration steps and performance metrics. The requirements.txt file was also updated to include new dependencies like datasets, gradio, matplotlib, and huggingface_hub.

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 a significant and valuable feature: fine-tuning capabilities for OmniGen2, covering both full-parameter and LoRA methods. The changes are extensive, including new training scripts, configuration examples, dataset handling logic, and updated documentation. The implementation is generally well-structured. However, I've identified a few critical issues in the new dataset sampling logic and callback implementations that must be addressed. There are also some medium-severity issues regarding documentation clarity and loss calculation for logging. After these fixes, this will be an excellent contribution to the project.