fix batch cfg (#93)

* fix batch cfg

* remove use_cfg

Author: akaitsuki-ii

diffsynth_engine/models/basic/transformer_helper.py

@@ -15,7 +15,7 @@ def __init__(self, dim: int, eps: float = 1e-6, device: str = "cuda:0", dtype: t
         self.silu = nn.SiLU()
 
     def forward(self, x, emb):
-        shift, scale = self.linear(self.silu(emb)).chunk(2, dim=1)
+        shift, scale = self.linear(self.silu(emb)).unsqueeze(1).chunk(2, dim=1)
         return modulate(self.norm(x), shift, scale)
 
 
@@ -27,7 +27,7 @@ def __init__(self, dim, device: str, dtype: torch.dtype):
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
 
     def forward(self, x, emb):
-        shift, scale, gate = self.linear(self.silu(emb)).chunk(3, dim=1)
+        shift, scale, gate = self.linear(self.silu(emb)).unsqueeze(1).chunk(3, dim=1)
         return modulate(self.norm(x), shift, scale), gate
 
 

diffsynth_engine/pipelines/flux_image.py

@@ -308,7 +308,6 @@ def __init__(
         vae_decoder: FluxVAEDecoder,
         vae_encoder: FluxVAEEncoder,
         load_text_encoder: bool = True,
-        use_cfg: bool = False,
         batch_cfg: bool = False,
         vae_tiled: bool = False,
         vae_tile_size: int = 256,
@@ -336,7 +335,6 @@ def __init__(
         self.vae_decoder = vae_decoder
         self.vae_encoder = vae_encoder
         self.load_text_encoder = load_text_encoder
-        self.use_cfg = use_cfg
         self.batch_cfg = batch_cfg
         self.ip_adapter = None
         self.redux = None
@@ -353,11 +351,15 @@ def __init__(
     def from_pretrained(
         cls,
         model_path_or_config: str | os.PathLike | FluxModelConfig,
+        load_text_encoder: bool = True,
+        batch_cfg: bool = False,
+        vae_tiled: bool = False,
+        vae_tile_size: int = 256,
+        vae_tile_stride: int = 256,
         control_type: ControlType = ControlType.normal,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.bfloat16,
         offload_mode: str | None = None,
-        load_text_encoder: bool = True,
         parallelism: int = 1,
         use_cfg_parallel: bool = False,
     ) -> "FluxImagePipeline":
@@ -454,6 +456,10 @@ def from_pretrained(
             vae_decoder=vae_decoder,
             vae_encoder=vae_encoder,
             load_text_encoder=load_text_encoder,
+            batch_cfg=batch_cfg,
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
             control_type=control_type,
             device=device,
             dtype=dtype,
@@ -530,10 +536,9 @@ def predict_noise_with_cfg(
         controlnet_params: List[ControlNetParams],
         current_step: int,
         total_step: int,
-        use_cfg: bool = False,
         batch_cfg: bool = False,
     ):
-        if cfg_scale <= 1.0 or not use_cfg:
+        if cfg_scale <= 1.0:
             return self.predict_noise(
                 latents,
                 timestep,
@@ -583,6 +588,10 @@ def predict_noise_with_cfg(
             add_text_embeds = torch.cat([positive_add_text_embeds, negative_add_text_embeds], dim=0)
             latents = torch.cat([latents, latents], dim=0)
             timestep = torch.cat([timestep, timestep], dim=0)
+            image_emb = torch.cat([image_emb, image_emb], dim=0) if image_emb is not None else None
+            image_ids = torch.cat([image_ids, image_ids], dim=0)
+            text_ids = torch.cat([text_ids, text_ids], dim=0)
+            guidance = torch.cat([guidance, guidance], dim=0)
             positive_noise_pred, negative_noise_pred = self.predict_noise(
                 latents,
                 timestep,
@@ -676,8 +685,14 @@ def prepare_latents(
                 num_inference_steps, mu=mu, sigma_min=1 / num_inference_steps, sigma_max=1.0
             )
             init_latents = latents.clone()
-        sigmas, timesteps = sigmas.to(device=self.device, dtype=self.dtype), timesteps.to(device=self.device, dtype=self.dtype)
-        init_latents, latents = init_latents.to(device=self.device, dtype=self.dtype), latents.to(device=self.device, dtype=self.dtype)
+        sigmas, timesteps = (
+            sigmas.to(device=self.device, dtype=self.dtype),
+            timesteps.to(device=self.device, dtype=self.dtype),
+        )
+        init_latents, latents = (
+            init_latents.to(device=self.device, dtype=self.dtype),
+            latents.to(device=self.device, dtype=self.dtype),
+        )
         return init_latents, latents, sigmas, timesteps
 
     def prepare_masked_latent(self, image: Image.Image, mask: Image.Image | None, height: int, width: int):
@@ -826,7 +841,7 @@ def __call__(
         # Encode prompts
         self.load_models_to_device(["text_encoder_1", "text_encoder_2"])
         positive_prompt_emb, positive_add_text_embeds = self.encode_prompt(prompt, clip_skip=clip_skip)
-        if self.use_cfg and cfg_scale > 1:
+        if cfg_scale > 1:
             negative_prompt_emb, negative_add_text_embeds = self.encode_prompt(negative_prompt, clip_skip=clip_skip)
         else:
             negative_prompt_emb, negative_add_text_embeds = None, None
@@ -868,7 +883,6 @@ def __call__(
                 controlnet_params=controlnet_params,
                 current_step=i,
                 total_step=len(timesteps),
-                use_cfg=self.use_cfg,
                 batch_cfg=self.batch_cfg,
             )
             # Denoise

diffsynth_engine/pipelines/sd_image.py

@@ -185,10 +185,13 @@ def __init__(
     def from_pretrained(
         cls,
         model_path_or_config: str | os.PathLike | SDModelConfig,
+        batch_cfg: bool = True,
+        vae_tiled: bool = False,
+        vae_tile_size: int = 256,
+        vae_tile_stride: int = 256,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.float16,
         offload_mode: str | None = None,
-        batch_cfg: bool = True,
     ) -> "SDImagePipeline":
         if isinstance(model_path_or_config, str):
             model_config = SDModelConfig(unet_path=model_path_or_config)
@@ -232,6 +235,9 @@ def from_pretrained(
             vae_decoder=vae_decoder,
             vae_encoder=vae_encoder,
             batch_cfg=batch_cfg,
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
             device=device,
             dtype=dtype,
         )
@@ -262,7 +268,7 @@ def predict_noise_with_cfg(
         cfg_scale: float,
         batch_cfg: bool = True,
     ):
-        if cfg_scale < 1.0:
+        if cfg_scale <= 1.0:
             return self.predict_noise(latents, timestep, positive_prompt_emb)
         if not batch_cfg:
             # cfg by predict noise one by one

diffsynth_engine/pipelines/sdxl_image.py

@@ -159,10 +159,13 @@ def __init__(
     def from_pretrained(
         cls,
         model_path_or_config: str | os.PathLike | SDXLModelConfig,
+        batch_cfg: bool = True,
+        vae_tiled: bool = False,
+        vae_tile_size: int = 256,
+        vae_tile_stride: int = 256,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.float16,
         offload_mode: str | None = None,
-        batch_cfg: bool = True,
     ) -> "SDXLImagePipeline":
         if isinstance(model_path_or_config, str):
             model_config = SDXLModelConfig(
@@ -220,6 +223,9 @@ def from_pretrained(
             vae_decoder=vae_decoder,
             vae_encoder=vae_encoder,
             batch_cfg=batch_cfg,
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
             device=device,
             dtype=dtype,
         )

diffsynth_engine/pipelines/wan_video.py

@@ -410,9 +410,13 @@ def __call__(
     def from_pretrained(
         cls,
         model_path_or_config: str | WanModelConfig,
+        shift: float | None = None,
+        batch_cfg: bool = False,
+        vae_tiled: bool = True,
+        vae_tile_size: Tuple[int, int] = (34, 34),
+        vae_tile_stride: Tuple[int, int] = (18, 16),
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
-        batch_cfg: bool = False,
         offload_mode: str | None = None,
         parallelism: int = 1,
         use_cfg_parallel: bool = False,
@@ -468,7 +472,7 @@ def from_pretrained(
             model_type = "1.3b-t2v"
 
         # shift for different model_type
-        shift = SHIFT_FACTORS[model_type]
+        shift = SHIFT_FACTORS[model_type] if shift is None else shift
 
         if parallelism > 1:
             parallel_config = cls.init_parallel_config(parallelism, use_cfg_parallel, model_config)
@@ -531,6 +535,9 @@ def from_pretrained(
             image_encoder=image_encoder,
             shift=shift,
             batch_cfg=batch_cfg,
+            vae_tiled=vae_tiled,
+            vae_tile_size=vae_tile_size,
+            vae_tile_stride=vae_tile_stride,
             device=device,
             dtype=dtype,
         )