[Diffusion] Add support for diffusion models (flagos-ai#767)

1) Implement DiffusionEngine, Adapter, Hook and Transform 
2) Add a simple LogIOTransform
3) Add example openjourney

---------

Co-authored-by: Yupu Feng <ypfeng@baai.ac.cn>
Co-authored-by: you-and-you <1823382186@qq.com>

Author: legitnull

.github/workflows/all-tests-nvidia.yml

@@ -25,8 +25,8 @@ jobs:
         run: |
           echo "ci_image=localhost:5000/flagscale:cuda12.8.1-cudnn9.7.1-python3.12-torch2.7.0-time2507111538" >> $GITHUB_OUTPUT  # Set output variable
 
-  # Megatron Unit Tests with Matrix
-  megatron-unit-tests:
+  # Train Megatron Unit Tests with Matrix
+  unit_tests_train_megatron:
     needs:
       - set-env
     uses: ./.github/workflows/unit-tests-nvidia.yml
@@ -47,34 +47,51 @@ jobs:
           - transformer/moe
           - transformer
           - ./
-    name: "megatron-${{ matrix.subset == './' && 'root' || matrix.subset }}"
+    name: "train_megatron-${{ matrix.subset == './' && 'root' || matrix.subset }}"
     with:
-      backend: megatron
+      backend: train_megatron
       subset: ${{ matrix.subset }}
       image: ${{ needs.set-env.outputs.ci_image }}
 
-  # Flagscale Unit Tests with Matrix
-  flagscale-unit-tests:
+  # Train Flagscale Unit Tests with Matrix
+  unit_tests_train_flagscale:
     needs:
       - set-env
-      - megatron-unit-tests
+      - unit_tests_train_megatron
     uses: ./.github/workflows/unit-tests-nvidia.yml
     strategy:
       matrix:
         subset:
           - runner
           - ./
-    name: "flagscale-${{ matrix.subset == './' && 'root' || matrix.subset }}"
+    name: "train_flagscale-${{ matrix.subset == './' && 'root' || matrix.subset }}"
     with:
-      backend: flagscale
+      backend: train_flagscale
+      subset: ${{ matrix.subset }}
+      image: ${{ needs.set-env.outputs.ci_image }}
+
+  # Inference Flagscale Unit Tests with Matrix
+  unit_tests_inference_flagscale:
+    needs:
+      - set-env
+      - unit_tests_train_flagscale
+    uses: ./.github/workflows/unit-tests-nvidia.yml
+    strategy:
+      matrix:
+        subset:
+          - inference
+          - transforms
+    name: "inference_flagscale-${{ matrix.subset == './' && 'root' || matrix.subset }}"
+    with:
+      backend: inference_flagscale
       subset: ${{ matrix.subset }}
       image: ${{ needs.set-env.outputs.ci_image }}
 
   # Functional Tests with Mision and Type Matrix
-  functional-tests-train:
+  functional_tests_train:
     needs:
       - set-env
-      - flagscale-unit-tests
+      - unit_tests_inference_flagscale
     uses: ./.github/workflows/functional-tests-nvidia.yml
     strategy:
       matrix:
@@ -89,10 +106,10 @@ jobs:
       type: train
       image: ${{ needs.set-env.outputs.ci_image }}
 
-  functional-tests-hetero:
+  functional_tests_hetero:
     needs:
       - set-env
-      - functional-tests-train
+      - functional_tests_train
     # TODO: test need fix
     uses: ./.github/workflows/functional-tests-nvidia.yml
     strategy:
@@ -105,10 +122,10 @@ jobs:
       type: hetero_train
       image: ${{ needs.set-env.outputs.ci_image }}
 
-  functional-tests-inference:
+  functional_tests_inference:
     needs:
       - set-env
-      - functional-tests-hetero
+      - functional_tests_hetero
     uses: ./.github/workflows/functional-tests-nvidia.yml
     strategy:
       matrix:
@@ -123,10 +140,10 @@ jobs:
       type: inference
       image: ${{ needs.set-env.outputs.ci_image }}
 
-  functional-tests-serve:
+  functional_tests_serve:
     needs:
       - set-env
-      - functional-tests-inference
+      - functional_tests_inference
     uses: ./.github/workflows/functional-tests-nvidia.yml
     strategy:
       matrix:
@@ -139,10 +156,10 @@ jobs:
       type: serve
       image: ${{ needs.set-env.outputs.ci_image }}
 
-  functional-tests-rl:
+  functional_tests_rl:
     needs:
       - set-env
-      - functional-tests-serve
+      - functional_tests_serve
     uses: ./.github/workflows/functional-tests-nvidia.yml
     strategy:
       matrix:
@@ -157,11 +174,12 @@ jobs:
   # Check All Tests
   all-tests:
     needs:
-      - megatron-unit-tests
-      - flagscale-unit-tests
-      - functional-tests-train
-      - functional-tests-hetero
-      - functional-tests-inference
+      - unit_tests_train_megatron
+      - unit_tests_train_flagscale
+      - unit_tests_inference_flagscale
+      - functional_tests_train
+      - functional_tests_hetero
+      - functional_tests_inference
       - functional-tests-serve
       - functional-tests-rl
     runs-on: ubuntu-latest

.github/workflows/unit-tests-nvidia.yml

@@ -53,10 +53,10 @@ jobs:
           pip install gitpython # TODO: Remove after updating the image
 
           git config --global --add safe.directory /__w/FlagScale/FlagScale
-          if [ "${{ inputs.backend }}" = "megatron" ] || [ "${{ inputs.backend }}" = "flagscale" ]; then
+          if [ "${{ inputs.backend }}" = "train_megatron" ] || [ "${{ inputs.backend }}" = "train_flagscale" ]; then
             echo ""
             # PYTHONPATH=./:$PYTHONPATH pip install . --config-settings=backend="Megatron-LM" --verbose --no-build-isolation
-          elif [ "${{ inputs.backend }}" = "vllm" ]; then
+          elif [ "${{ inputs.backend }}" = "vllm" ] || [ "${{ inputs.backend }}" = "inference_flagscale" ]; then
             source /root/miniconda3/bin/activate flagscale-inference
             pip install scikit-build scikit-build-core
             pip install --no-build-isolation git+https://github.com/FlagOpen/FlagGems.git@release_v1.0.0

examples/openjourney/conf/inference.yaml

@@ -0,0 +1,26 @@
+defaults:
+  - _self_
+  - inference: t2i
+
+experiment:
+  exp_name: openjourney
+  exp_dir: outputs/${experiment.exp_name}
+  model: prompthero/openjourney
+  task:
+    type: inference
+    backend: vllm # TODO(yupu): Remove this restriction
+    entrypoint: flagscale/inference/diffusion_entrypoint.py
+  runner:
+    hostfile: null
+  cmds:
+    before_start: source /root/miniconda3/bin/activate flagscale-inference
+  envs:
+    CUDA_VISIBLE_DEVICES: 0
+    CUDA_DEVICE_MAX_CONNECTIONS: 1
+    # Optionally, set HF_HOME and HF_ENDPOINT
+
+action: run
+
+hydra:
+  run:
+    dir: ${experiment.exp_dir}/hydra

examples/openjourney/conf/inference/t2i.yaml

@@ -0,0 +1,18 @@
+engine:
+  model_config:
+    model: ${experiment.model}
+    loader: diffusers
+  engine_config:
+    results_path: ${experiment.exp_dir}/results
+    transformations:
+      LogIOTransformation:
+        log_level: "info"
+      StateScopeTransformation:
+        {}
+
+generate:
+  prompt: "retro serie of different cars with different colors and shapes, mdjrny-v4 style"
+  width: 512
+  height: 512
+  num_inference_steps: 50
+  guidance_scale: 7.0

flagscale/inference/diffusion_entrypoint.py

@@ -0,0 +1,46 @@
+import argparse
+
+from typing import Union
+
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from flagscale.inference.inference_engine import InferenceEngine
+from flagscale.runner.utils import logger
+
+
+def parse_config() -> Union[DictConfig, ListConfig]:
+    """Parse the configuration file"""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config-path", type=str, required=True, help="Path to the configuration YAML file"
+    )
+    args = parser.parse_args()
+
+    config = OmegaConf.load(args.config_path)
+
+    # TODO(yupu): Any checks?
+
+    return config
+
+
+def inference(cfg: DictConfig) -> None:
+    """Run the model inference
+
+    Args:
+        cfg: The parsed configuration
+    """
+
+    engine = InferenceEngine(cfg.get("engine", {}))
+
+    generate_cfg = cfg.get("generate", {})
+    outputs = engine.generate(**generate_cfg)
+
+    engine.save(outputs)
+
+
+if __name__ == "__main__":
+    parsed_cfg = parse_config()
+    assert isinstance(parsed_cfg, DictConfig)  # To make pyright happy
+    inference(parsed_cfg)
+    logger.info("Inference model inference completed")