mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
Qwen-Image-Edit
This commit is contained in:
Haoming
GitHub
parent
6bbcf3a787
commit
b2226faacc
@ -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)
|
||||
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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 <b>Flux-Kontext</b> Only<br>
|
||||
Use in <b>txt2img</b> to achieve the effect of EmptySD3LatentImage with custom resolution
|
||||
For <b>Flux-Kontext</b> and <b>Qwen-Image-Edit</b><br>
|
||||
Use in <b>txt2img</b> to achieve the effect of empty latent with custom resolution
|
||||
"""
|
||||
|
||||
i2i_info = """
|
||||
For <b>Flux-Kontext</b> Only<br>
|
||||
For <b>Flux-Kontext</b> and <b>Qwen-Image-Edit</b><br>
|
||||
Use in <b>img2img</b> to achieve the effect of 2 input images<br>
|
||||
<b>NOTE:</b> 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
|
||||
|
||||
@ -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 ; <b>Note:</b> 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"),
|
||||
|
||||
Loading…
Reference in New Issue
Block a user