theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-21 16:18:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

qwen-image splited training

Browse Source
This commit is contained in:
Artiprocher
2025-09-02 16:44:14 +08:00
parent 260e32217f
commit b6da77e468
7 changed files with 221 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
diffsynth/pipelines/qwen_image.py
View File

@@ -174,9 +174,12 @@ class QwenImagePipeline(BasePipeline):
computation_dtype=self.torch_dtype,
computation_device="cuda",
)
enable_vram_management(self.text_encoder, module_map=module_map, module_config=model_config)
enable_vram_management(self.dit, module_map=module_map, module_config=model_config)
enable_vram_management(self.vae, module_map=module_map, module_config=model_config)
if self.text_encoder is not None:
enable_vram_management(self.text_encoder, module_map=module_map, module_config=model_config)
if self.dit is not None:
enable_vram_management(self.dit, module_map=module_map, module_config=model_config)
if self.vae is not None:
enable_vram_management(self.vae, module_map=module_map, module_config=model_config)
def enable_vram_management(self, num_persistent_param_in_dit=None, vram_limit=None, vram_buffer=0.5, enable_dit_fp8_computation=False):