diffsynth-engine 0.4.2.dev6__tar.gz → 0.4.3.dev1__tar.gz

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.2.dev6
+Version: 0.4.3.dev1
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/diffsynth_engine/configs/pipeline.py

@@ -16,6 +16,7 @@ class BaseConfig:
     vae_tile_stride: int | Tuple[int, int] = 256
     device: str = "cuda"
     offload_mode: Optional[str] = None
+    offload_to_disk: bool = False
 
 
 @dataclass
@@ -62,11 +63,13 @@ class SDPipelineConfig(BaseConfig):
         model_path: str | os.PathLike | List[str | os.PathLike],
         device: str = "cuda",
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "SDPipelineConfig":
         return cls(
             model_path=model_path,
             device=device,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
 
@@ -87,11 +90,13 @@ class SDXLPipelineConfig(BaseConfig):
         model_path: str | os.PathLike | List[str | os.PathLike],
         device: str = "cuda",
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "SDXLPipelineConfig":
         return cls(
             model_path=model_path,
             device=device,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
 
@@ -116,6 +121,7 @@ class FluxPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Ba
         device: str = "cuda",
         parallelism: int = 1,
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "FluxPipelineConfig":
         return cls(
             model_path=model_path,
@@ -123,6 +129,7 @@ class FluxPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Ba
             parallelism=parallelism,
             use_fsdp=True,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
     def __post_init__(self):
@@ -160,6 +167,7 @@ class WanPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Bas
         device: str = "cuda",
         parallelism: int = 1,
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "WanPipelineConfig":
         return cls(
             model_path=model_path,
@@ -169,6 +177,7 @@ class WanPipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, Bas
             use_cfg_parallel=True,
             use_fsdp=True,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
     def __post_init__(self):
@@ -196,6 +205,7 @@ class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfi
         device: str = "cuda",
         parallelism: int = 1,
         offload_mode: Optional[str] = None,
+        offload_to_disk: bool = False,
     ) -> "QwenImagePipelineConfig":
         return cls(
             model_path=model_path,
@@ -206,6 +216,7 @@ class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfi
             use_cfg_parallel=True,
             use_fsdp=True,
             offload_mode=offload_mode,
+            offload_to_disk=offload_to_disk,
         )
 
     def __post_init__(self):

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/diffsynth_engine/models/qwen_image/qwen_image_dit.py

@@ -48,8 +48,8 @@ class QwenEmbedRope(nn.Module):
         super().__init__()
         self.theta = theta
         self.axes_dim = axes_dim
-        pos_index = torch.arange(1024)
-        neg_index = torch.arange(1024).flip(0) * -1 - 1
+        pos_index = torch.arange(10000)
+        neg_index = torch.arange(10000).flip(0) * -1 - 1
         self.pos_freqs = torch.cat(
             [
                 self.rope_params(pos_index, self.axes_dim[0], self.theta),

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/diffsynth_engine/pipelines/base.py

@@ -41,6 +41,7 @@ class BasePipeline:
         self.offload_mode = None
         self.model_names = []
         self._offload_param_dict = {}
+        self.offload_to_disk = False
 
     @classmethod
     def from_pretrained(cls, model_path_or_config: str | BaseConfig) -> "BasePipeline":
@@ -228,19 +229,23 @@ class BasePipeline:
                 model.eval()
         return self
 
-    def enable_cpu_offload(self, offload_mode: str):
-        valid_offload_mode = ("cpu_offload", "sequential_cpu_offload")
+    def enable_cpu_offload(self, offload_mode: str | None, offload_to_disk:bool = False):
+        valid_offload_mode = ("cpu_offload", "sequential_cpu_offload", "disable", None)
         if offload_mode not in valid_offload_mode:
             raise ValueError(f"offload_mode must be one of {valid_offload_mode}, but got {offload_mode}")
         if self.device == "cpu" or self.device == "mps":
             logger.warning("must set an non cpu device for pipeline before calling enable_cpu_offload")
             return
-        if offload_mode == "cpu_offload":
+        if offload_mode is None or offload_mode == "disable":
+            self._disable_offload()
+        elif offload_mode == "cpu_offload":
             self._enable_model_cpu_offload()
         elif offload_mode == "sequential_cpu_offload":
             self._enable_sequential_cpu_offload()
+        self.offload_to_disk = offload_to_disk
 
-    def _enable_model_cpu_offload(self):
+
+    def _enable_model_cpu_offload(self):        
         for model_name in self.model_names:
             model = getattr(self, model_name)
             if model is not None:
@@ -253,6 +258,15 @@ class BasePipeline:
             if model is not None:
                 enable_sequential_cpu_offload(model, self.device)
         self.offload_mode = "sequential_cpu_offload"
+    
+    def _disable_offload(self):
+        self.offload_mode = None                
+        self._offload_param_dict = {}        
+        for model_name in self.model_names:
+            model = getattr(self, model_name)
+            if model is not None:
+                model.to(self.device)
+
 
     def enable_fp8_autocast(
         self, model_names: List[str], compute_dtype: torch.dtype = torch.bfloat16, use_fp8_linear: bool = False
@@ -260,6 +274,7 @@ class BasePipeline:
         for model_name in model_names:
             model = getattr(self, model_name)
             if model is not None:
+                model.to(device=self.device, dtype=torch.float8_e4m3fn)                
                 enable_fp8_autocast(model, compute_dtype, use_fp8_linear)
         self.fp8_autocast_enabled = True
 
@@ -282,10 +297,26 @@ class BasePipeline:
         # load the needed models to device
         for model_name in load_model_names:
             model = getattr(self, model_name)
+            if model is None:
+                raise ValueError(f"model {model_name} is not loaded, maybe this model has been destroyed by model_lifecycle_finish function with offload_to_disk=True")
             if model is not None and (p := next(model.parameters(), None)) is not None and p.device.type != self.device:
                 model.to(self.device)
         # fresh the cuda cache
         empty_cache()
 
+    def model_lifecycle_finish(self, model_names: List[str] | None = None):
+        if not self.offload_to_disk or self.offload_mode is None:
+            return 
+        for model_name in model_names:
+            model = getattr(self, model_name)
+            del model
+            if model_name in self._offload_param_dict:
+                del self._offload_param_dict[model_name]
+            setattr(self, model_name, None)
+            print(f"model {model_name} has been deleted from memory")
+            logger.info(f"model {model_name} has been deleted from memory")
+            empty_cache()
+        
+        
     def compile(self):
         raise NotImplementedError(f"{self.__class__.__name__} does not support compile")

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/diffsynth_engine/pipelines/qwen_image.py

@@ -41,19 +41,32 @@ class QwenImageLoRAConverter(LoRAStateDictConverter):
         dit_dict = {}
         for key, param in lora_state_dict.items():
             origin_key = key
-            if "lora_A.default.weight" not in key:
+            lora_a_suffix = None
+            if "lora_A.default.weight" in key:
+                lora_a_suffix = "lora_A.default.weight"
+            elif "lora_A.weight" in key:
+                lora_a_suffix = "lora_A.weight"
+
+            if lora_a_suffix is None:
                 continue
+
             lora_args = {}
             lora_args["down"] = param
-            lora_args["up"] = lora_state_dict[origin_key.replace("lora_A.default.weight", "lora_B.default.weight")]
+
+            lora_b_suffix = lora_a_suffix.replace("lora_A", "lora_B")
+            lora_args["up"] = lora_state_dict[origin_key.replace(lora_a_suffix, lora_b_suffix)]
+
             lora_args["rank"] = lora_args["up"].shape[1]
-            alpha_key = origin_key.replace("lora_A.default.weight", "alpha").replace("lora_up.default.weight", "alpha")
+            alpha_key = origin_key.replace("lora_up", "lora_A").replace(lora_a_suffix, "alpha")
+
             if alpha_key in lora_state_dict:
                 alpha = lora_state_dict[alpha_key]
             else:
                 alpha = lora_args["rank"]
             lora_args["alpha"] = alpha
-            key = key.replace(".lora_A.default.weight", "")
+
+            key = key.replace(f".{lora_a_suffix}", "")
+
             if key.startswith("transformer") and "attn.to_out.0" in key:
                 key = key.replace("attn.to_out.0", "attn.to_out")
             dit_dict[key] = lora_args
@@ -195,7 +208,19 @@ class QwenImagePipeline(BasePipeline):
         pipe.eval()
 
         if config.offload_mode is not None:
-            pipe.enable_cpu_offload(config.offload_mode)
+            pipe.enable_cpu_offload(config.offload_mode, config.offload_to_disk)
+        
+        if config.model_dtype == torch.float8_e4m3fn:
+            pipe.dtype = torch.bfloat16  # compute dtype
+            pipe.enable_fp8_autocast(
+                model_names=["dit"], compute_dtype=pipe.dtype, use_fp8_linear=config.use_fp8_linear
+            )
+
+        if config.encoder_dtype == torch.float8_e4m3fn:
+            pipe.dtype = torch.bfloat16  # compute dtype
+            pipe.enable_fp8_autocast(
+                model_names=["encoder"], compute_dtype=pipe.dtype, use_fp8_linear=config.use_fp8_linear
+            )
 
         if config.parallelism > 1:
             pipe = ParallelWrapper(
@@ -380,6 +405,7 @@ class QwenImagePipeline(BasePipeline):
             negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(negative_prompt, 1, 4096)
         else:
             negative_prompt_embeds, negative_prompt_embeds_mask = None, None
+        self.model_lifecycle_finish(["encoder"])
 
         hide_progress = dist.is_initialized() and dist.get_rank() != 0
         for i, timestep in enumerate(tqdm(timesteps, disable=hide_progress)):
@@ -399,6 +425,7 @@ class QwenImagePipeline(BasePipeline):
             # UI
             if progress_callback is not None:
                 progress_callback(i, len(timesteps), "DENOISING")
+        self.model_lifecycle_finish(["dit"])
         # Decode image
         self.load_models_to_device(["vae"])
         latents = rearrange(latents, "B C H W -> B C 1 H W")
@@ -410,5 +437,6 @@ class QwenImagePipeline(BasePipeline):
         )
         image = self.vae_output_to_image(vae_output)
         # Offload all models
+        self.model_lifecycle_finish(["vae"])        
         self.load_models_to_device([])
         return image

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/diffsynth_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.2.dev6
+Version: 0.4.3.dev1
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/docs/tutorial.md

@@ -88,7 +88,6 @@ We will continuously update DiffSynth-Engine to support more models. (Wan2.2 LoR
 
 After the model is downloaded, load the model with the corresponding pipeline and perform inference.
 
-
 ### Image Generation(Qwen-Image)
 
 The following code calls `QwenImagePipeline` to load the [Qwen-Image](https://www.modelscope.cn/models/Qwen/Qwen-Image) model and generate an image. Recommended resolutions are 928×1664, 1104×1472, 1328×1328, 1472×1104, and 1664×928, with a suggested cfg_scale of 4. If no negative_prompt is provided, it defaults to a single space character (not an empty string). For multi-GPU parallelism, currently only cfg parallelism is supported (parallelism=2), with other optimization efforts underway.
@@ -122,7 +121,7 @@ image.save("image.png")
 
 Please note that if some necessary modules, like text encoders, are missing from a model repository, the pipeline will automatically download the required files.
 
-#### Detailed Parameters(Qwen-Image)
+### Detailed Parameters(Qwen-Image)
 
 In the image generation pipeline `pipe`, we can use the following parameters for fine-grained control:
 

{diffsynth_engine-0.4.2.dev6 → diffsynth_engine-0.4.3.dev1}/docs/tutorial_zh.md

@@ -2,13 +2,13 @@
 
 ## 安装
 
-在使用 DiffSynth-Engine 前，请先确保您的硬件设备满足以下要求：
+在使用 DiffSynth-Engine 前，请先确保您的硬件设备满足以下要求:
 
 * NVIDIA GPU CUDA 计算能力 8.6+（例如 RTX 50 Series、RTX 40 Series、RTX 30 Series 等，详见 [NVidia 文档](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities)）或 Apple Silicon M 系列芯片
 
-以及 Python 环境需求：Python 3.10+。
+以及 Python 环境需求: Python 3.10+。
 
-使用 `pip3` 工具从 PyPI 安装 DiffSynth-Engine：
+使用 `pip3` 工具从 PyPI 安装 DiffSynth-Engine:
 
 ```shell
 pip3 install diffsynth-engine
@@ -64,7 +64,7 @@ model_path = fetch_model("Wan-AI/Wan2.1-T2V-14B", path="diffusion_pytorch_model*
 
 ## 模型类型
 
-Diffusion 模型包含多种多样的模型结构，每种模型由对应的流水线进行加载和推理，目前我们支持的模型类型包括：
+Diffusion 模型包含多种多样的模型结构，每种模型由对应的流水线进行加载和推理，目前我们支持的模型类型包括:
 
 | 模型结构         | 样例                                                         | 流水线              |
 | --------------- | ------------------------------------------------------------ | ------------------- |
@@ -123,16 +123,17 @@ image.save("image.png")
 
 #### 详细参数(Qwen-Image)
 
-在图像生成流水线 `pipe` 中，我们可以通过以下参数进行精细的控制：
+在图像生成流水线 `pipe` 中，我们可以通过以下参数进行精细的控制:
 
 * `prompt`: 提示词，用于描述生成图像的内容，支持多种语言(中文/英文/日文等)，例如“一只猫”/"a cat"/"庭を走る猫"。
 * `negative_prompt`: 负面提示词，用于描述不希望图像中出现的内容，例如“ugly”，默认为一个空格而不是空字符串， " "。
-* `cfg_scale`: [Classifier-free guidance](https://arxiv.org/abs/2207.12598) 的引导系数，通常更大的引导系数可以达到更强的文图相关性，但会降低生成内容的多样性，推荐值为4。
+* `cfg_scale`:[Classifier-free guidance](https://arxiv.org/abs/2207.12598) 的引导系数，通常更大的引导系数可以达到更强的文图相关性，但会降低生成内容的多样性，推荐值为4。
 * `height`: 图像高度。
 * `width`: 图像宽度。
 * `num_inference_steps`: 推理步数，通常推理步数越多，计算时间越长，图像质量越高。
 * `seed`: 随机种子，固定的随机种子可以使生成的内容固定。
 
+
 ### 图像生成
 
 以下代码可以调用 `FluxImagePipeline` 加载[麦橘超然](https://www.modelscope.cn/models/MAILAND/majicflus_v1/summary?version=v1.0)模型生成一张图。如果要加载其他结构的模型，请将代码中的 `FluxImagePipeline` 和 `FluxPipelineConfig` 替换成对应的流水线模块及配置。
@@ -152,7 +153,7 @@ image.save("image.png")
 
 #### 详细参数
 
-在图像生成流水线 `pipe` 中，我们可以通过以下参数进行精细的控制：
+在图像生成流水线 `pipe` 中，我们可以通过以下参数进行精细的控制:
 
 * `prompt`: 提示词，用于描述生成图像的内容，例如“a cat”。
 * `negative_prompt`: 负面提示词，用于描述不希望图像中出现的内容，例如“ugly”。
@@ -217,7 +218,7 @@ save_video(video, "video.mp4")
 
 #### 详细参数
 
-在视频生成流水线 `pipe` 中，我们可以通过以下参数进行精细的控制：
+在视频生成流水线 `pipe` 中，我们可以通过以下参数进行精细的控制:
 
 * `prompt`: 提示词，用于描述生成图像的内容，例如“a cat”。
 * `negative_prompt`: 负面提示词，用于描述不希望图像中出现的内容，例如“ugly”。