add fluxcontrolpipeline support

examples/flux_control_example.py

@@ -0,0 +1,104 @@
+import logging
+import time
+import torch
+
+from transformers import T5EncoderModel
+from xfuser import xFuserArgs
+from xfuser.config import FlexibleArgumentParser
+from xfuser.core.distributed import (
+    get_world_group,
+    get_data_parallel_rank,
+    get_data_parallel_world_size,
+    get_runtime_state,
+    get_data_parallel_world_size,
+)
+from xfuser.model_executor.pipelines.pipeline_flux_control import xFuserPipelineFluxControlPipeline
+from diffusers.utils import load_image  
+
+def main():
+    parser = FlexibleArgumentParser(description="xFuser Arguments")
+    args = xFuserArgs.add_cli_args(parser).parse_args()
+    engine_args = xFuserArgs.from_cli_args(args)
+    engine_config, input_config = engine_args.create_config()
+    engine_config.runtime_config.dtype = torch.bfloat16
+    local_rank = get_world_group().local_rank
+    text_encoder_2 = T5EncoderModel.from_pretrained(engine_config.model_config.model, subfolder="text_encoder_2", torch_dtype=torch.bfloat16)
+
+    if args.use_fp8_t5_encoder:
+        from optimum.quanto import freeze, qfloat8, quantize
+        logging.info(f"rank {local_rank} quantizing text encoder 2")
+        quantize(text_encoder_2, weights=qfloat8)
+        freeze(text_encoder_2)
+
+    cache_args = {
+            "use_teacache": engine_args.use_teacache,
+            "use_fbcache": engine_args.use_fbcache,
+            "rel_l1_thresh": 0.12,
+            "return_hidden_states_first": False,
+            "num_steps": input_config.num_inference_steps,
+        }
+
+    pipe = xFuserPipelineFluxControlPipeline.from_pretrained(
+        pretrained_model_name_or_path=engine_config.model_config.model,
+        engine_config=engine_config,
+        cache_args=cache_args,
+        torch_dtype=torch.bfloat16,
+        text_encoder_2=text_encoder_2,
+    )
+
+    if args.enable_sequential_cpu_offload:
+        pipe.enable_sequential_cpu_offload(gpu_id=local_rank)
+        logging.info(f"rank {local_rank} sequential CPU offload enabled")
+    else:
+        pipe = pipe.to(f"cuda:{local_rank}")
+
+    parameter_peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}")
+
+    pipe.prepare_run(input_config, steps=input_config.num_inference_steps)
+
+    # load control_image
+    control_image = load_image("test_imgs/016#.png").convert("RGB").resize((1024,1024))
+
+    torch.cuda.reset_peak_memory_stats()
+    start_time = time.time()
+    output = pipe(
+        height=input_config.height,
+        width=input_config.width,
+        prompt=input_config.prompt,
+        control_image=control_image,
+        num_inference_steps=input_config.num_inference_steps,
+        output_type=input_config.output_type,
+        max_sequence_length=input_config.max_sequence_length,
+        guidance_scale=input_config.guidance_scale,
+        generator=torch.Generator(device="cuda").manual_seed(input_config.seed),
+    )
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}")
+
+    parallel_info = (
+        f"dp{engine_args.data_parallel_degree}_cfg{engine_config.parallel_config.cfg_degree}_"
+        f"ulysses{engine_args.ulysses_degree}_ring{engine_args.ring_degree}_"
+        f"tp{engine_args.tensor_parallel_degree}_"
+        f"pp{engine_args.pipefusion_parallel_degree}_patch{engine_args.num_pipeline_patch}"
+    )
+    if input_config.output_type == "pil":
+        dp_group_index = get_data_parallel_rank()
+        num_dp_groups = get_data_parallel_world_size()
+        dp_batch_size = (input_config.batch_size + num_dp_groups - 1) // num_dp_groups
+        if pipe.is_dp_last_group():
+            for i, image in enumerate(output.images):
+                image_rank = dp_group_index * dp_batch_size + i
+                image_name = f"flux_control_result_{parallel_info}_{image_rank}_tc_{engine_args.use_torch_compile}.png"
+                image.save(f"./results/{image_name}")
+                print(f"image {i} saved to ./results/{image_name}")
+
+    if get_world_group().rank == get_world_group().world_size - 1:
+        print(
+            f"epoch time: {elapsed_time:.2f} sec, parameter memory: {parameter_peak_memory/1e9:.2f} GB, memory: {peak_memory/1e9:.2f} GB"
+        )
+    get_runtime_state().destroy_distributed_env()
+
+
+if __name__ == "__main__":
+    main()

examples/run.sh

@@ -11,6 +11,7 @@ declare -A MODEL_CONFIGS=(
     ["Pixart-sigma"]="pixartsigma_example.py /cfs/dit/PixArt-Sigma-XL-2-2K-MS 20"
     ["Sd3"]="sd3_example.py /cfs/dit/stable-diffusion-3-medium-diffusers 20"
     ["Flux"]="flux_example.py /cfs/dit/FLUX.1-dev/ 28"
+    ["FluxControl"]="flux_control_example.py /cfs/dit/FLUX.1-Depth-dev/ 28"
     ["HunyuanDiT"]="hunyuandit_example.py /cfs/dit/HunyuanDiT-v1.2-Diffusers 50"
     ["SDXL"]="sdxl_example.py /cfs/dit/stable-diffusion-xl-base-1.0 30"
 )