diff --git a/docs/en/Pipeline_Usage/GPU_support.md b/docs/en/Pipeline_Usage/GPU_support.md index aba5706..a60c119 100644 --- a/docs/en/Pipeline_Usage/GPU_support.md +++ b/docs/en/Pipeline_Usage/GPU_support.md @@ -81,4 +81,5 @@ Set 0 or not set: indicates not enabling the binding function #### Parameters for specific models | Model | Parameter | Note | |----------------|---------------------------|-------------------| -| Wan 14B series | --initialize_model_on_cpu | The 14B model needs to be initialized on the CPU | \ No newline at end of file +| Wan 14B series | --initialize_model_on_cpu | The 14B model needs to be initialized on the CPU | +| Qwen-Image series | --initialize_model_on_cpu | The model needs to be initialized on the CPU | \ No newline at end of file diff --git a/docs/zh/Pipeline_Usage/GPU_support.md b/docs/zh/Pipeline_Usage/GPU_support.md index 8124147..b130fad 100644 --- a/docs/zh/Pipeline_Usage/GPU_support.md +++ b/docs/zh/Pipeline_Usage/GPU_support.md @@ -81,4 +81,5 @@ export CPU_AFFINITY_CONF=1 #### 特定模型需要开启的参数 | 模型 | 参数 | 备注 | |-----------|------|-------------------| -| Wan 14B系列 | --initialize_model_on_cpu | 14B模型需要在cpu上进行初始化 | \ No newline at end of file +| Wan 14B系列 | --initialize_model_on_cpu | 14B模型需要在cpu上进行初始化 | +| Qwen-Image系列 | --initialize_model_on_cpu | 模型需要在cpu上进行初始化 | \ No newline at end of file diff --git a/examples/qwen_image/model_training/special/npu_training/Qwen-Image-Edit-2509-NPU.sh b/examples/qwen_image/model_training/special/npu_training/Qwen-Image-Edit-2509-NPU.sh index 4892228..02de9e9 100644 --- a/examples/qwen_image/model_training/special/npu_training/Qwen-Image-Edit-2509-NPU.sh +++ b/examples/qwen_image/model_training/special/npu_training/Qwen-Image-Edit-2509-NPU.sh @@ -16,4 +16,5 @@ accelerate launch --config_file examples/qwen_image/model_training/full/accelera --output_path "./models/train/Qwen-Image-Edit-2509_full" \ --trainable_models "dit" \ --use_gradient_checkpointing \ - --find_unused_parameters + --find_unused_parameters \ + --initialize_model_on_cpu diff --git a/examples/qwen_image/model_training/train.py b/examples/qwen_image/model_training/train.py index 8f38d04..ecb4239 100644 --- a/examples/qwen_image/model_training/train.py +++ b/examples/qwen_image/model_training/train.py @@ -101,6 +101,7 @@ def qwen_image_parser(): parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.") parser.add_argument("--processor_path", type=str, default=None, help="Path to the processor. If provided, the processor will be used for image editing.") parser.add_argument("--zero_cond_t", default=False, action="store_true", help="A special parameter introduced by Qwen-Image-Edit-2511. Please enable it for this model.") + parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.") return parser @@ -151,7 +152,7 @@ if __name__ == "__main__": fp8_models=args.fp8_models, offload_models=args.offload_models, task=args.task, - device=accelerator.device, + device="cpu" if args.initialize_model_on_cpu else accelerator.device, zero_cond_t=args.zero_cond_t, ) model_logger = ModelLogger(