fix merge

Author: zhtmike

examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py

@@ -85,11 +85,9 @@ def log_validation(pipeline, args, accelerator, generator, global_step, is_final
         os.makedirs(val_save_dir)
 
     original_image = (
-        lambda image_url_or_path: (
-            load_image(image_url_or_path)
-            if urlparse(image_url_or_path).scheme
-            else Image.open(image_url_or_path).convert("RGB")
-        )
+        lambda image_url_or_path: load_image(image_url_or_path)
+        if urlparse(image_url_or_path).scheme
+        else Image.open(image_url_or_path).convert("RGB")
     )(args.val_image_url_or_path)
 
     if torch.backends.mps.is_available():

src/diffusers/loaders/peft.py

@@ -46,7 +46,7 @@
 logger = logging.get_logger(__name__)
 
 _SET_ADAPTER_SCALE_FN_MAPPING = defaultdict(
-    lambda: lambda model_cls, weights: weights,
+    lambda: (lambda model_cls, weights: weights),
     {
         "UNet2DConditionModel": _maybe_expand_lora_scales,
         "UNetMotionModel": _maybe_expand_lora_scales,

src/diffusers/models/attention_dispatch.py

@@ -3212,29 +3212,26 @@ def _flash_attention_3_varlen_hub(
     if _parallel_config is not None and _parallel_config.context_parallel_config.ring_degree > 1:
         raise NotImplementedError("`ring_degree > 1` is not yet supported for the _FLASH_3_VARLEN_HUB backend.")
 
-    lse = None
     batch_size, seq_len_q, _, _ = query.shape
     _, seq_len_kv, _, _ = key.shape
 
     if _parallel_config is None:
         if attn_mask is not None:
             attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
-
-        (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
-            _prepare_for_flash_attn_or_sage_varlen(
-                batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+            (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+                _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device)
             )
-        )
-
-        key_valid, value_valid = [], []
-        for b in range(batch_size):
-            valid_len = seqlens_k[b]
-            key_valid.append(key[b, :valid_len])
-            value_valid.append(value[b, :valid_len])
+            indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten()
+            key_packed = key.reshape(-1, *key.shape[2:])[indices_k]
+            value_packed = value.reshape(-1, *value.shape[2:])[indices_k]
+        else:
+            (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+                _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device)
+            )
+            key_packed = key.flatten(0, 1)
+            value_packed = value.flatten(0, 1)
 
         query_packed = query.flatten(0, 1)
-        key_packed = torch.cat(key_valid, dim=0)
-        value_packed = torch.cat(value_valid, dim=0)
 
         func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_VARLEN_HUB].kernel_fn
         result = func(