From b2226faacc0a3d716b0ce8046cdd935224a99992 Mon Sep 17 00:00:00 2001
From: Haoming <73768377+Haoming02@users.noreply.github.com>
Date: Sat, 18 Oct 2025 15:11:01 +0800
Subject: [PATCH] Qwen-Image-Edit
---
backend/diffusion_engine/qwen.py | 53 ++++++++-
backend/loader.py | 1 +
backend/nn/flux.py | 2 +-
backend/nn/qwen.py | 6 +-
backend/nn/svdq.py | 7 +-
backend/text_processing/qwen_engine.py | 112 +++++++++++++-----
.../scripts/image_stitch.py | 14 ++-
modules/shared_options.py | 2 +-
8 files changed, 154 insertions(+), 43 deletions(-)
diff --git a/backend/diffusion_engine/qwen.py b/backend/diffusion_engine/qwen.py
index 3b8d225c..53ee9457 100644
--- a/backend/diffusion_engine/qwen.py
+++ b/backend/diffusion_engine/qwen.py
@@ -1,7 +1,13 @@
+import math
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from modules.prompt_parser import SdConditioning
+
import torch
from huggingface_guess import model_list
-from backend import memory_management
+from backend import args, memory_management
from backend.diffusion_engine.base import ForgeDiffusionEngine, ForgeObjects
from backend.modules.k_prediction import PredictionDiscreteFlow
from backend.patcher.clip import CLIP
@@ -37,22 +43,63 @@ class QwenImage(ForgeDiffusionEngine):
self.is_wan = True
+ self.images_vl = []
+ self.ref_latents = []
+ self.image_prompt = ""
+
def set_clip_skip(self, clip_skip):
pass
@torch.inference_mode()
- def get_learned_conditioning(self, prompt: list[str]):
+ def get_learned_conditioning(self, prompt: "SdConditioning"):
memory_management.load_model_gpu(self.forge_objects.clip.patcher)
+ if not prompt.is_negative_prompt and self.image_prompt:
+ return self.get_learned_conditioning_with_image(prompt)
return self.text_processing_engine_qwen(prompt)
+ @torch.inference_mode()
+ def get_learned_conditioning_with_image(self, prompt: list[str]):
+ cond = self.text_processing_engine_qwen([self.image_prompt + "".join(prompt)], images=self.images_vl)
+ args.dynamic_args["ref_latents"] = self.ref_latents.copy()
+ self.images_vl.clear()
+ self.ref_latents.clear()
+ self.image_prompt = ""
+ return cond
+
@torch.inference_mode()
def get_prompt_lengths_on_ui(self, prompt):
token_count = len(self.text_processing_engine_qwen.tokenize([prompt])[0])
return token_count, max(999, token_count)
+ @torch.inference_mode()
+ def encode_vision(self, image):
+ samples = image.movedim(-1, 1) # b, c, h, w
+
+ total = int(384 * 384)
+ scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
+ width = round(samples.shape[3] * scale_by)
+ height = round(samples.shape[2] * scale_by)
+
+ s = torch.nn.functional.interpolate(samples, size=(height, width), mode="area")
+ self.images_vl.append(s.movedim(1, -1))
+
+ total = int(1024 * 1024)
+ scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
+ width = round(samples.shape[3] * scale_by / 8.0) * 8
+ height = round(samples.shape[2] * scale_by / 8.0) * 8
+
+ s = torch.nn.functional.interpolate(samples, size=(height, width), mode="area")
+ self.ref_latents.append(self.forge_objects.vae.encode(s.movedim(1, -1)[:, :, :, :3]))
+
+ self.image_prompt += f"Picture {len(self.images_vl)}: <|vision_start|><|image_pad|><|vision_end|>"
+
@torch.inference_mode()
def encode_first_stage(self, x):
- sample = self.forge_objects.vae.encode(x.movedim(2, -1) * 0.5 + 0.5)
+ if x.size(0) > 1:
+ x = x[0].unsqueeze(0) # enforce batch_size of 1
+ start_image = x.movedim(1, -1) * 0.5 + 0.5
+ self.encode_vision(start_image)
+ sample = self.forge_objects.vae.encode(start_image)
sample = self.forge_objects.vae.first_stage_model.process_in(sample)
return sample.to(x)
diff --git a/backend/loader.py b/backend/loader.py
index 4e4345ca..d93ebaac 100644
--- a/backend/loader.py
+++ b/backend/loader.py
@@ -566,6 +566,7 @@ def forge_loader(sd: os.PathLike, additional_state_dicts: list[os.PathLike] = No
repo_name = estimated_config.huggingface_repo
backend.args.dynamic_args["kontext"] = "kontext" in str(sd).lower()
+ backend.args.dynamic_args["edit"] = "qwen" in str(sd).lower() and "edit" in str(sd).lower()
backend.args.dynamic_args["nunchaku"] = getattr(estimated_config, "nunchaku", False)
if getattr(estimated_config, "nunchaku", False):
diff --git a/backend/nn/flux.py b/backend/nn/flux.py
index da108e13..26f990f1 100644
--- a/backend/nn/flux.py
+++ b/backend/nn/flux.py
@@ -424,7 +424,7 @@ class IntegratedFluxTransformer2DModel(nn.Module):
else:
h_offset = h
- kontext, kontext_ids = process_img(ref, index=1, h_offset=h_offset, w_offset=w_offset)
+ kontext, kontext_ids = process_img(ref.to(x), index=1, h_offset=h_offset, w_offset=w_offset)
img = torch.cat([img, kontext], dim=1)
img_ids = torch.cat([img_ids, kontext_ids], dim=1)
h = max(h, ref.shape[-2] + h_offset)
diff --git a/backend/nn/qwen.py b/backend/nn/qwen.py
index 32c02844..d291b3c1 100644
--- a/backend/nn/qwen.py
+++ b/backend/nn/qwen.py
@@ -15,6 +15,7 @@ if xformers_enabled():
else:
from backend.attention import attention_pytorch as attention_function
+from backend.args import dynamic_args
from backend.nn.flux import EmbedND
from backend.utils import pad_to_patch_size
@@ -400,6 +401,9 @@ class QwenImageTransformer2DModel(nn.Module):
hidden_states, img_ids, orig_shape = self.process_img(x)
num_embeds = hidden_states.shape[1]
+ if dynamic_args.get("ref_latents", None) is not None:
+ ref_latents = dynamic_args["ref_latents"]
+
if ref_latents is not None:
h = 0
w = 0
@@ -421,7 +425,7 @@ class QwenImageTransformer2DModel(nn.Module):
h = max(h, ref.shape[-2] + h_offset)
w = max(w, ref.shape[-1] + w_offset)
- kontext, kontext_ids, _ = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset)
+ kontext, kontext_ids, _ = self.process_img(ref.to(x), index=index, h_offset=h_offset, w_offset=w_offset)
hidden_states = torch.cat([hidden_states, kontext], dim=1)
img_ids = torch.cat([img_ids, kontext_ids], dim=1)
diff --git a/backend/nn/svdq.py b/backend/nn/svdq.py
index bdc4d202..8e4694b7 100644
--- a/backend/nn/svdq.py
+++ b/backend/nn/svdq.py
@@ -100,7 +100,7 @@ class SVDQFluxTransformer2DModel(nn.Module):
else:
h_offset = h
- kontext, kontext_ids = process_img(ref, index=1, h_offset=h_offset, w_offset=w_offset)
+ kontext, kontext_ids = process_img(ref.to(x), index=1, h_offset=h_offset, w_offset=w_offset)
img = torch.cat([img, kontext], dim=1)
img_ids = torch.cat([img_ids, kontext_ids], dim=1)
h = max(h, ref.shape[-2] + h_offset)
@@ -602,6 +602,9 @@ class NunchakuQwenImageTransformer2DModel(NunchakuModelMixin, QwenImageTransform
hidden_states, img_ids, orig_shape = self.process_img(x)
num_embeds = hidden_states.shape[1]
+ if dynamic_args.get("ref_latents", None) is not None:
+ ref_latents = dynamic_args["ref_latents"]
+
if ref_latents is not None:
h = 0
w = 0
@@ -623,7 +626,7 @@ class NunchakuQwenImageTransformer2DModel(NunchakuModelMixin, QwenImageTransform
h = max(h, ref.shape[-2] + h_offset)
w = max(w, ref.shape[-1] + w_offset)
- kontext, kontext_ids, _ = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset)
+ kontext, kontext_ids, _ = self.process_img(ref.to(x), index=index, h_offset=h_offset, w_offset=w_offset)
hidden_states = torch.cat([hidden_states, kontext], dim=1)
img_ids = torch.cat([img_ids, kontext_ids], dim=1)
diff --git a/backend/text_processing/qwen_engine.py b/backend/text_processing/qwen_engine.py
index 532d3bcf..a7a3bb5c 100644
--- a/backend/text_processing/qwen_engine.py
+++ b/backend/text_processing/qwen_engine.py
@@ -22,23 +22,23 @@ class QwenTextProcessingEngine:
self.max_length = 99999999
self.min_length = 1
self.id_pad = 151643
+ self.id_template = 151644
+ self.id_image = 151655
self.llama_template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+ self.image_template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
- def tokenize(self, texts):
- llama_texts = [self.llama_template.format(text) for text in texts]
+ def tokenize(self, texts, template=None):
+ llama_texts = [(template or self.llama_template).format(text) for text in texts]
return self.tokenizer(llama_texts)["input_ids"]
- def encode_with_transformers(self, tokens):
- device = memory_management.text_encoder_device()
- tokens = tokens.to(device)
- self.text_encoder.to(device=device)
- return self.text_encoder(x=tokens)
-
- def tokenize_line(self, line):
+ def tokenize_line(self, line, images=None):
parsed = parsing.parse_prompt_attention(line, self.emphasis.name)
- tokenized = self.tokenize([text for text, _ in parsed])
+ tokenized = self.tokenize(
+ [text for text, _ in parsed],
+ self.image_template if bool(images) else self.llama_template,
+ )
chunks = []
chunk = PromptChunk()
@@ -64,9 +64,15 @@ class QwenTextProcessingEngine:
next_chunk()
continue
+ embed_count = 0
position = 0
while position < len(tokens):
token = tokens[position]
+
+ if token == self.id_image:
+ token = {"type": "image", "data": images[embed_count], "original_type": "image"}
+ embed_count += 1
+
chunk.tokens.append(token)
chunk.multipliers.append(weight)
position += 1
@@ -76,7 +82,7 @@ class QwenTextProcessingEngine:
return chunks, token_count
- def __call__(self, texts):
+ def __call__(self, texts, images=None):
zs = []
cache = {}
@@ -86,7 +92,7 @@ class QwenTextProcessingEngine:
if line in cache:
line_z_values = cache[line]
else:
- chunks, token_count = self.tokenize_line(line)
+ chunks, token_count = self.tokenize_line(line, images)
line_z_values = []
# pad all chunks to length of longest chunk
@@ -104,7 +110,7 @@ class QwenTextProcessingEngine:
multipliers += [1.0] * remaining_count
z = self.process_tokens([tokens], [multipliers])[0]
- z = self.postprocess_tokens(z, tokens)
+ z = self.strip_template(z, tokens)
line_z_values.append(z)
cache[line] = line_z_values
@@ -112,16 +118,18 @@ class QwenTextProcessingEngine:
return torch.stack(zs)
- def postprocess_tokens(self, out, tokens):
- """strip the llama_template"""
+ def strip_template(self, out, tokens):
template_end = 0
count_im_start = 0
for i, v in enumerate(tokens):
- elem = int(v)
- if elem == 151644 and count_im_start < 2:
- template_end = i
- count_im_start += 1
+ try:
+ elem = int(v)
+ if elem == self.id_template and count_im_start < 2:
+ template_end = i
+ count_im_start += 1
+ except TypeError:
+ continue
if out.shape[1] > (template_end + 3):
if int(tokens[template_end + 1]) == 872:
@@ -130,15 +138,61 @@ class QwenTextProcessingEngine:
return out[template_end:]
+ def process_embeds(self, batch_tokens):
+ device = memory_management.text_encoder_device()
+
+ embeds_out = []
+ attention_masks = []
+ num_tokens = []
+
+ for tokens in batch_tokens:
+ attention_mask = []
+ tokens_temp = []
+ other_embeds = []
+ eos = False
+ index = 0
+
+ for t in tokens:
+ try:
+ token = int(t)
+ attention_mask.append(0 if eos else 1)
+ tokens_temp += [token]
+ if not eos and token == self.id_pad:
+ eos = True
+ except TypeError:
+ other_embeds.append((index, t))
+ index += 1
+
+ tokens_embed = torch.tensor([tokens_temp], device=device, dtype=torch.long)
+ tokens_embed = self.text_encoder.get_input_embeddings()(tokens_embed)
+
+ index = 0
+ embeds_info = []
+
+ for o in other_embeds:
+ emb, extra = self.text_encoder.preprocess_embed(o[1], device=device)
+ if emb is None:
+ index += -1
+ continue
+
+ ind = index + o[0]
+ emb = emb.view(1, -1, emb.shape[-1]).to(device=device, dtype=torch.float32)
+ emb_shape = emb.shape[1]
+
+ assert emb.shape[-1] == tokens_embed.shape[-1]
+ tokens_embed = torch.cat([tokens_embed[:, :ind], emb, tokens_embed[:, ind:]], dim=1)
+ attention_mask = attention_mask[:ind] + [1] * emb_shape + attention_mask[ind:]
+ index += emb_shape - 1
+ emb_type = o[1].get("type", None)
+ embeds_info.append({"type": emb_type, "index": ind, "size": emb_shape, "extra": extra})
+
+ embeds_out.append(tokens_embed)
+ attention_masks.append(attention_mask)
+ num_tokens.append(sum(attention_mask))
+
+ return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens, embeds_info
+
def process_tokens(self, batch_tokens, batch_multipliers):
- tokens = torch.asarray(batch_tokens)
-
- z, _ = self.encode_with_transformers(tokens)
-
- self.emphasis.tokens = batch_tokens
- self.emphasis.multipliers = torch.asarray(batch_multipliers).to(z)
- self.emphasis.z = z
- self.emphasis.after_transformers()
- z = self.emphasis.z
-
+ embeds, mask, count, info = self.process_embeds(batch_tokens)
+ z, _ = self.text_encoder(x=None, embeds=embeds, attention_mask=mask, num_tokens=count, embeds_info=info)
return z
diff --git a/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py b/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
index 290fa53b..233b0c94 100644
--- a/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
+++ b/extensions-builtin/sd_forge_image_stitch/scripts/image_stitch.py
@@ -15,12 +15,12 @@ from modules.sd_samplers_common import approximation_indexes, images_tensor_to_s
from modules.shared import device, opts
t2i_info = """
-For Flux-Kontext Only
-Use in txt2img to achieve the effect of EmptySD3LatentImage with custom resolution
+For Flux-Kontext and Qwen-Image-Edit
+Use in txt2img to achieve the effect of empty latent with custom resolution
"""
i2i_info = """
-For Flux-Kontext Only
+For Flux-Kontext and Qwen-Image-Edit
Use in img2img to achieve the effect of 2 input images
NOTE: This doesn't actually stitch the images, so use "1st/2nd" instead of "left/right" in prompts
"""
@@ -58,8 +58,7 @@ class ImageStitch(scripts.Script):
def process(self, p: "StableDiffusionProcessing", reference: "Image.Image"):
if reference is None:
return
- if not dynamic_args.get("kontext", False):
- # print("\nImageStitch only works with Flux-Kontext!\n")
+ if not any(dynamic_args[key] for key in ("kontext", "edit")):
return
image = images.flatten(reference, opts.img2img_background_color)
@@ -67,8 +66,11 @@ class ImageStitch(scripts.Script):
image = np.moveaxis(image, 2, 0)
image = torch.from_numpy(image).to(device=device, dtype=torch.float32)
- dynamic_args["ref_latents"] = images_tensor_to_samples(
+ ref = images_tensor_to_samples(
image.unsqueeze(0),
approximation_indexes.get(opts.sd_vae_encode_method),
p.sd_model,
)
+
+ if dynamic_args["kontext"]:
+ dynamic_args["ref_latents"] = ref
diff --git a/modules/shared_options.py b/modules/shared_options.py
index 568b0aa2..09421262 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -281,7 +281,7 @@ options_templates.update(
("optimizations", "Optimizations", "sd"),
{
"cross_attention_optimization": OptionInfo("Automatic", "Cross Attention Optimization", gr.Dropdown, {"choices": ("Automatic",), "interactive": False}),
- "persistent_cond_cache": OptionInfo(True, "Persistent Cond Cache").info("do not recalculate conds if the prompts and parameters have not changed since previous generation"),
+ "persistent_cond_cache": OptionInfo(False, "Persistent Cond Cache").info("do not re-encode prompts if only the Seed changes ; Note: breaks Qwen-Image-Edit if only the input image was changed"),
"skip_early_cond": OptionInfo(0.0, "Ignore Negative Prompt during Early Steps", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}, infotext="Skip Early CFG").info("in percentage of total steps; 0 = disable; higher = faster"),
"s_min_uncond": OptionInfo(0.0, "Skip Negative Prompt during Later Steps", gr.Slider, {"minimum": 0.0, "maximum": 8.0, "step": 0.05}).info('in "sigma"; 0 = disable; higher = faster'),
"s_min_uncond_all": OptionInfo(False, "For the above option, skip every step", infotext="NGMS all steps").info("otherwise, only skip every other step"),