From 9b06cfef808b284cd6e696ac042fde4f24c8d141 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sat, 5 Sep 2026 13:36:18 +0200
Subject: [PATCH] add llada editing

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 modules/processing_args.py                 |  5 +++++
 pipelines/llada/transformer_llada_image.py | 10 ++++++++--
 pipelines/model_llada.py                   |  1 -
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/modules/processing_args.py b/modules/processing_args.py
index 52a68cdce..42088f8b9 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -110,6 +110,11 @@ def task_specific_kwargs(p, model):
                 'height': p.height,
                 'input_images': [p.init_images], # omnigen expects list-of-lists
             }
+        elif model_cls == 'LLaDAImagePipeline':
+            task_args = {
+                'generation_mode': 'editing',
+                'image': p.init_images[0],
+            }
     elif task_type == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images', [])) > 0:
         p.ops.append('instruct')
         task_args = {
diff --git a/pipelines/llada/transformer_llada_image.py b/pipelines/llada/transformer_llada_image.py
index d1aa783d2..29cb06d00 100644
--- a/pipelines/llada/transformer_llada_image.py
+++ b/pipelines/llada/transformer_llada_image.py
@@ -1705,7 +1705,10 @@ class LLaDAImageSigVQEmbeddings(nn.Module):
         hidden_size = position_embedding.shape[1]
         original_size = int(position_embedding.shape[0] ** 0.5)
         position_embedding = position_embedding.reshape(original_size, original_size, hidden_size)
-        position_embedding = position_embedding.permute(2, 0, 1).unsqueeze(0).float()
+        position_embedding = position_embedding.permute(2, 0, 1).unsqueeze(0).to(
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
 
         height_coordinates = torch.arange(grid_height, device=hidden_states.device, dtype=torch.float32)
         width_coordinates = torch.arange(grid_width, device=hidden_states.device, dtype=torch.float32)
@@ -1738,7 +1741,10 @@ class LLaDAImageSigVQQuantizer(nn.Module):
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states.permute(0, 2, 3, 1).contiguous()
         hidden_states = F.normalize(hidden_states.reshape(-1, hidden_states.shape[-1]), p=2, dim=-1)
-        embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+        embedding = F.normalize(self.embedding.weight, p=2, dim=-1).to(
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
         distances = (
             torch.sum(hidden_states**2, dim=1, keepdim=True)
             + torch.sum(embedding**2, dim=1)
diff --git a/pipelines/model_llada.py b/pipelines/model_llada.py
index 647cd059c..bd0a1297e 100644
--- a/pipelines/model_llada.py
+++ b/pipelines/model_llada.py
@@ -69,7 +69,6 @@ def load_llada_image(checkpoint_info, diffusers_load_config=None):
     )
     pipe.task_args = {
         'output_type': 'np',
-        'generation_mode': 'text',
     }
     # generation_mode = "text", "vq", "editing"