import torch, math from PIL import Image from typing import Union from tqdm import tqdm from einops import rearrange import numpy as np from math import prod import cv2 from torchvision.transforms import v2, InterpolationMode from ..diffusion import FlowMatchScheduler from ..core import ModelConfig, gradient_checkpoint_forward from ..diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput from ..models.qwen_image_dit import QwenImageDiT from ..models.qwen_image_text_encoder import QwenImageTextEncoder from ..models.qwen_image_vae import QwenImageVAE from ..models.qwen_image_controlnet import QwenImageBlockWiseControlNet from ..models.qwen_image_connector import QwenImageConnector from ..models.qwen_image_learnable_query import QwenImageLearnableQuery from ..models.qwen_image_extractor import QwenImageExtractor from ..models.qwen_image_rformer import RFormer,RFormer2D from ..models.qwen_image_mae import QwenImageMAE import torch import random torch.set_printoptions(threshold=10_000_000) def extract_polygons(item): """ 递归解析 edit_region，提取所有的多边形轮廓。 目标是找到所有的 list of points [[x,y], [x,y], ...] """ if not isinstance(item, list): return [] if len(item) == 2 and all(isinstance(x, (int, float)) for x in item): return None first_elem = item[0] if len(item) > 0 else None if isinstance(first_elem, list) and len(first_elem) == 2 and all(isinstance(x, (int, float)) for x in first_elem): return [np.array(item, dtype=np.float32)] polys = [] for sub in item: res = extract_polygons(sub) if res: polys.extend(res) return polys class QwenImagePipeline(BasePipeline): def __init__(self, device="cuda", torch_dtype=torch.bfloat16): super().__init__( device=device, torch_dtype=torch_dtype, height_division_factor=16, width_division_factor=16, ) from transformers import Qwen2Tokenizer, Qwen2VLProcessor self.scheduler = FlowMatchScheduler("Qwen-Image") self.text_encoder: QwenImageTextEncoder = None self.connector: QwenImageConnector=None self.learnable_query: QwenImageLearnableQuery=None self.extractor: QwenImageExtractor=None self.rformer: RFormer=None self.rformer2d: RFormer2D=None self.mae: QwenImageMAE=None self.dit: QwenImageDiT = None self.vae: QwenImageVAE = None self.blockwise_controlnet: QwenImageBlockwiseMultiControlNet = None self.tokenizer: Qwen2Tokenizer = None self.processor: Qwen2VLProcessor = None self.in_iteration_models = ("dit", "blockwise_controlnet") self.units = [ QwenImageUnit_ShapeChecker(), QwenImageUnit_NoiseInitializer(), QwenImageUnit_InputImageEmbedder(), QwenImageUnit_Inpaint(), QwenImageUnit_PromptEmbedderWithExtractor(), QwenImageUnit_MirrorPPROperationEncoder(), QwenImageUnit_MirrorPPROperationEncoder2D(), QwenImageUnit_EditImageEmbedder(), QwenImageUnit_ContextImageEmbedder(), QwenImageUnit_PromptEmbedder(), QwenImageUnit_EntityControl(), QwenImageUnit_BlockwiseControlNet(), ] self.model_fn = model_fn_qwen_image @staticmethod def from_pretrained( torch_dtype: torch.dtype = torch.bfloat16, device: Union[str, torch.device] = "cuda", model_configs: list[ModelConfig] = [], tokenizer_config: ModelConfig = ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"), processor_config: ModelConfig = None, vram_limit: float = None, ): pipe = QwenImagePipeline(device=device, torch_dtype=torch_dtype) model_pool = pipe.download_and_load_models(model_configs, vram_limit) pipe.text_encoder = model_pool.fetch_model("qwen_image_text_encoder") pipe.dit = model_pool.fetch_model("qwen_image_dit") pipe.vae = model_pool.fetch_model("qwen_image_vae") pipe.blockwise_controlnet = QwenImageBlockwiseMultiControlNet(model_pool.fetch_model("qwen_image_blockwise_controlnet", index="all")) pipe.connector=model_pool.fetch_model("qwen_image_connector") pipe.learnable_query=model_pool.fetch_model("qwen_image_learnable_query") pipe.extractor=model_pool.fetch_model("qwen_image_extractor") pipe.rformer=model_pool.fetch_model("qwen_image_rformer") pipe.rformer2d=model_pool.fetch_model("qwen_image_rformer2D") pipe.mae=model_pool.fetch_model("qwen_image_mae") if tokenizer_config is not None: tokenizer_config.download_if_necessary() from transformers import Qwen2Tokenizer pipe.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_config.path) if processor_config is not None: processor_config.download_if_necessary() from transformers import Qwen2VLProcessor pipe.processor = Qwen2VLProcessor.from_pretrained(processor_config.path) pipe.vram_management_enabled = pipe.check_vram_management_state() if pipe.vram_management_enabled: print("enable vram management") else: print("disable vram management") return pipe @torch.no_grad() def __call__( self, prompt: str, negative_prompt: str = "", cfg_scale: float = 4.0, input_image: Image.Image = None, denoising_strength: float = 1.0, inpaint_mask: Image.Image = None, inpaint_blur_size: int = None, inpaint_blur_sigma: float = None, height: int = 1328, width: int = 1328, seed: int = None, rand_device: str = "cpu", num_inference_steps: int = 30, exponential_shift_mu: float = None, blockwise_controlnet_inputs: list[ControlNetInput] = None, eligen_entity_prompts: list[str] = None, eligen_entity_masks: list[Image.Image] = None, eligen_enable_on_negative: bool = False, edit_image: Image.Image = None, edit_image_auto_resize: bool = True, edit_rope_interpolation: bool = False, zero_cond_t: bool = False, example_origin: Image.Image = None, example_target: Image.Image = None, shape_1=None, no_qwenvl=False, context_image: Image.Image = None, tiled: bool = False, tile_size: int = 128, tile_stride: int = 64, progress_bar_cmd = tqdm, ): self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, dynamic_shift_len=(height // 16) * (width // 16), exponential_shift_mu=exponential_shift_mu) inputs_posi = { "prompt": prompt, } inputs_nega = { "negative_prompt": negative_prompt, } inputs_shared = { "cfg_scale": cfg_scale, "input_image": input_image, "denoising_strength": denoising_strength, "inpaint_mask": inpaint_mask, "inpaint_blur_size": inpaint_blur_size, "inpaint_blur_sigma": inpaint_blur_sigma, "height": height, "width": width, "seed": seed, "rand_device": rand_device, "num_inference_steps": num_inference_steps, "blockwise_controlnet_inputs": blockwise_controlnet_inputs, "tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride, "eligen_entity_prompts": eligen_entity_prompts, "eligen_entity_masks": eligen_entity_masks, "eligen_enable_on_negative": eligen_enable_on_negative, "edit_image": edit_image, "edit_image_auto_resize": edit_image_auto_resize, "edit_rope_interpolation": edit_rope_interpolation, "context_image": context_image, "zero_cond_t": zero_cond_t, "example_origin": example_origin, "example_target": example_target,"shape_1":shape_1,"no_qwenvl":no_qwenvl } for unit in self.units: inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) self.load_models_to_device(self.in_iteration_models) models = {name: getattr(self, name) for name in self.in_iteration_models} for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) noise_pred = self.cfg_guided_model_fn( self.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) self.load_models_to_device(['vae']) image = self.vae.decode(inputs_shared["latents"], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) image = self.vae_output_to_image(image) self.load_models_to_device([]) return image class QwenImageBlockwiseMultiControlNet(torch.nn.Module): def __init__(self, models: list[QwenImageBlockWiseControlNet]): super().__init__() if not isinstance(models, list): models = [models] self.models = torch.nn.ModuleList(models) for model in models: if hasattr(model, "vram_management_enabled") and getattr(model, "vram_management_enabled"): self.vram_management_enabled = True def preprocess(self, controlnet_inputs: list[ControlNetInput], conditionings: list[torch.Tensor], **kwargs): processed_conditionings = [] for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): conditioning = rearrange(conditioning, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2) model_output = self.models[controlnet_input.controlnet_id].process_controlnet_conditioning(conditioning) processed_conditionings.append(model_output) return processed_conditionings def blockwise_forward(self, image, conditionings: list[torch.Tensor], controlnet_inputs: list[ControlNetInput], progress_id, num_inference_steps, block_id, **kwargs): res = 0 for controlnet_input, conditioning in zip(controlnet_inputs, conditionings): progress = (num_inference_steps - 1 - progress_id) / max(num_inference_steps - 1, 1) if progress > controlnet_input.start + (1e-4) or progress < controlnet_input.end - (1e-4): continue model_output = self.models[controlnet_input.controlnet_id].blockwise_forward(image, conditioning, block_id) res = res + model_output * controlnet_input.scale return res class QwenImageUnit_Inpaint(PipelineUnit): def __init__(self): super().__init__( input_params=("inpaint_mask", "height", "width", "inpaint_blur_size", "inpaint_blur_sigma"), output_params=("inpaint_mask",), ) def process(self, pipe: QwenImagePipeline, inpaint_mask, height, width, inpaint_blur_size, inpaint_blur_sigma): if inpaint_mask is None: return {} inpaint_mask = pipe.preprocess_image(inpaint_mask.convert("RGB").resize((width // 8, height // 8)), min_value=0, max_value=1) inpaint_mask = inpaint_mask.mean(dim=1, keepdim=True) if inpaint_blur_size is not None and inpaint_blur_sigma is not None: from torchvision.transforms import GaussianBlur blur = GaussianBlur(kernel_size=inpaint_blur_size * 2 + 1, sigma=inpaint_blur_sigma) inpaint_mask = blur(inpaint_mask) return {"inpaint_mask": inpaint_mask} class QwenImageUnit_EntityControl(PipelineUnit): def __init__(self): super().__init__( take_over=True, input_params=("eligen_entity_prompts", "width", "height", "eligen_enable_on_negative", "cfg_scale"), output_params=("entity_prompt_emb", "entity_masks", "entity_prompt_emb_mask"), onload_model_names=("text_encoder",) ) def extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() valid_lengths = bool_mask.sum(dim=1) selected = hidden_states[bool_mask] split_result = torch.split(selected, valid_lengths.tolist(), dim=0) return split_result def get_prompt_emb(self, pipe: QwenImagePipeline, prompt) -> dict: if pipe.text_encoder is not None: prompt = [prompt] 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" drop_idx = 34 txt = [template.format(e) for e in prompt] txt_tokens = pipe.tokenizer(txt, max_length=1024+drop_idx, padding=True, truncation=True, return_tensors="pt").to(pipe.device) hidden_states = pipe.text_encoder(input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, output_hidden_states=True,)[-1] split_hidden_states = self.extract_masked_hidden(hidden_states, txt_tokens.attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max([e.size(0) for e in split_hidden_states]) prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) prompt_embeds = prompt_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) return {"prompt_emb": prompt_embeds, "prompt_emb_mask": encoder_attention_mask} else: return {} def preprocess_masks(self, pipe, masks, height, width, dim): out_masks = [] for mask in masks: mask = pipe.preprocess_image(mask.resize((width, height), resample=Image.NEAREST)).mean(dim=1, keepdim=True) > 0 mask = mask.repeat(1, dim, 1, 1).to(device=pipe.device, dtype=pipe.torch_dtype) out_masks.append(mask) return out_masks def prepare_entity_inputs(self, pipe, entity_prompts, entity_masks, width, height): entity_masks = self.preprocess_masks(pipe, entity_masks, height//8, width//8, 1) entity_masks = torch.cat(entity_masks, dim=0).unsqueeze(0) prompt_embs, prompt_emb_masks = [], [] for entity_prompt in entity_prompts: prompt_emb_dict = self.get_prompt_emb(pipe, entity_prompt) prompt_embs.append(prompt_emb_dict['prompt_emb']) prompt_emb_masks.append(prompt_emb_dict['prompt_emb_mask']) return prompt_embs, prompt_emb_masks, entity_masks def prepare_eligen(self, pipe, prompt_emb_nega, eligen_entity_prompts, eligen_entity_masks, width, height, enable_eligen_on_negative, cfg_scale): entity_prompt_emb_posi, entity_prompt_emb_posi_mask, entity_masks_posi = self.prepare_entity_inputs(pipe, eligen_entity_prompts, eligen_entity_masks, width, height) if enable_eligen_on_negative and cfg_scale != 1.0: entity_prompt_emb_nega = [prompt_emb_nega['prompt_emb']] * len(entity_prompt_emb_posi) entity_prompt_emb_nega_mask = [prompt_emb_nega['prompt_emb_mask']] * len(entity_prompt_emb_posi) entity_masks_nega = entity_masks_posi else: entity_prompt_emb_nega, entity_prompt_emb_nega_mask, entity_masks_nega = None, None, None eligen_kwargs_posi = {"entity_prompt_emb": entity_prompt_emb_posi, "entity_masks": entity_masks_posi, "entity_prompt_emb_mask": entity_prompt_emb_posi_mask} eligen_kwargs_nega = {"entity_prompt_emb": entity_prompt_emb_nega, "entity_masks": entity_masks_nega, "entity_prompt_emb_mask": entity_prompt_emb_nega_mask} return eligen_kwargs_posi, eligen_kwargs_nega def process(self, pipe: QwenImagePipeline, inputs_shared, inputs_posi, inputs_nega): eligen_entity_prompts, eligen_entity_masks = inputs_shared.get("eligen_entity_prompts", None), inputs_shared.get("eligen_entity_masks", None) if eligen_entity_prompts is None or eligen_entity_masks is None or len(eligen_entity_prompts) == 0 or len(eligen_entity_masks) == 0: return inputs_shared, inputs_posi, inputs_nega pipe.load_models_to_device(self.onload_model_names) eligen_enable_on_negative = inputs_shared.get("eligen_enable_on_negative", False) eligen_kwargs_posi, eligen_kwargs_nega = self.prepare_eligen(pipe, inputs_nega, eligen_entity_prompts, eligen_entity_masks, inputs_shared["width"], inputs_shared["height"], eligen_enable_on_negative, inputs_shared["cfg_scale"]) inputs_posi.update(eligen_kwargs_posi) if inputs_shared.get("cfg_scale", 1.0) != 1.0: inputs_nega.update(eligen_kwargs_nega) return inputs_shared, inputs_posi, inputs_nega class QwenImageUnit_BlockwiseControlNet(PipelineUnit): def __init__(self): super().__init__( input_params=("blockwise_controlnet_inputs", "tiled", "tile_size", "tile_stride"), output_params=("blockwise_controlnet_conditioning",), onload_model_names=("vae",) ) def apply_controlnet_mask_on_latents(self, pipe, latents, mask): mask = (pipe.preprocess_image(mask) + 1) / 2 mask = mask.mean(dim=1, keepdim=True) mask = 1 - torch.nn.functional.interpolate(mask, size=latents.shape[-2:]) latents = torch.concat([latents, mask], dim=1) return latents def apply_controlnet_mask_on_image(self, pipe, image, mask): mask = mask.resize(image.size) mask = pipe.preprocess_image(mask).mean(dim=[0, 1]).cpu() image = np.array(image) image[mask > 0] = 0 image = Image.fromarray(image) return image def process(self, pipe: QwenImagePipeline, blockwise_controlnet_inputs: list[ControlNetInput], tiled, tile_size, tile_stride): if blockwise_controlnet_inputs is None: return {} pipe.load_models_to_device(self.onload_model_names) conditionings = [] for controlnet_input in blockwise_controlnet_inputs: image = controlnet_input.image if controlnet_input.inpaint_mask is not None: image = self.apply_controlnet_mask_on_image(pipe, image, controlnet_input.inpaint_mask) image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) image = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) if controlnet_input.inpaint_mask is not None: image = self.apply_controlnet_mask_on_latents(pipe, image, controlnet_input.inpaint_mask) conditionings.append(image) return {"blockwise_controlnet_conditioning": conditionings} class QwenImageUnit_ContextImageEmbedder(PipelineUnit): def __init__(self): super().__init__( input_params=("context_image", "height", "width", "tiled", "tile_size", "tile_stride"), output_params=("context_latents",), onload_model_names=("vae",) ) def process(self, pipe: QwenImagePipeline, context_image, height, width, tiled, tile_size, tile_stride): if context_image is None: return {} pipe.load_models_to_device(self.onload_model_names) context_image = pipe.preprocess_image(context_image.resize((width, height))).to(device=pipe.device, dtype=pipe.torch_dtype) context_latents = pipe.vae.encode(context_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) return {"context_latents": context_latents} class QwenImageUnit_ShapeChecker(PipelineUnit): def __init__(self): super().__init__( input_params=("height", "width"), output_params=("height", "width"), ) def process(self, pipe: QwenImagePipeline, height, width): height, width = pipe.check_resize_height_width(height, width) return {"height": height, "width": width} class QwenImageUnit_NoiseInitializer(PipelineUnit): def __init__(self): super().__init__( input_params=("height", "width", "seed", "rand_device"), output_params=("noise",), ) def process(self, pipe: QwenImagePipeline, height, width, seed, rand_device): noise = pipe.generate_noise((1, 16, height//8, width//8), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype) return {"noise": noise} class QwenImageUnit_InputImageEmbedder(PipelineUnit): def __init__(self): super().__init__( input_params=("input_image", "noise", "tiled", "tile_size", "tile_stride","edit_region","loss_weight_mask"), output_params=("latents", "input_latents","loss_weight_mask"), onload_model_names=("vae",) ) def process(self, pipe: QwenImagePipeline, input_image, noise, tiled, tile_size, tile_stride, edit_region=None,loss_weight_mask=None): if input_image is None: return {"latents": noise, "input_latents": None, "loss_weight_mask": None} pipe.load_models_to_device(['vae']) image = pipe.preprocess_image(input_image).to(device=pipe.device, dtype=pipe.torch_dtype) input_latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) latent_mask = None if edit_region is not None: B, C, H_lat, W_lat = input_latents.shape latent_mask = torch.ones((B, 1, H_lat, W_lat), device=input_latents.device, dtype=input_latents.dtype) edit_weight = 100.0 scale_h = H_lat / image.shape[2] scale_w = W_lat / image.shape[3] for b in range(B): if B == 1: current_region = edit_region else: try: current_region = edit_region[b] except: current_region = edit_region if not current_region: continue polygons = extract_polygons(current_region) if len(polygons) > 0: mask_np = np.zeros((H_lat, W_lat), dtype=np.float32) scaled_polys = [] for poly in polygons: poly_scaled = poly * np.array([scale_w, scale_h], dtype=np.float32) scaled_polys.append(poly_scaled.astype(np.int32)) cv2.fillPoly(mask_np, scaled_polys, color=1.0) mask_tensor = torch.from_numpy(mask_np).to(latent_mask.device) latent_mask[b, 0] = torch.where(mask_tensor > 0.5, torch.tensor(edit_weight, device=latent_mask.device), latent_mask[b, 0]) if loss_weight_mask is not None: latent_mask=loss_weight_mask if pipe.scheduler.training: return { "latents": noise, "input_latents": input_latents, "loss_weight_mask": latent_mask } else: latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) return { "latents": latents, "input_latents": input_latents } class QwenImageUnit_PromptEmbedder(PipelineUnit): def __init__(self): super().__init__( seperate_cfg=True, input_params_posi={"prompt": "prompt"}, input_params_nega={"prompt": "negative_prompt"}, input_params=("edit_image",), output_params=("prompt_emb", "prompt_emb_mask"), onload_model_names=("text_encoder",) ) def extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() valid_lengths = bool_mask.sum(dim=1) selected = hidden_states[bool_mask] split_result = torch.split(selected, valid_lengths.tolist(), dim=0) return split_result def calculate_dimensions(self, target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height def resize_image(self, image, target_area=384*384): width, height = self.calculate_dimensions(target_area, image.size[0] / image.size[1]) return image.resize((width, height)) def encode_prompt(self, pipe: QwenImagePipeline, prompt): 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" drop_idx = 34 txt = [template.format(e) for e in prompt] model_inputs = pipe.tokenizer(txt, max_length=4096+drop_idx, padding=True, truncation=True, return_tensors="pt").to(pipe.device) if model_inputs.input_ids.shape[1] >= 1024: print(f"Warning!!! QwenImage model was trained on prompts up to 512 tokens. Current prompt requires {model_inputs['input_ids'].shape[1] - drop_idx} tokens, which may lead to unpredictable behavior.") hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, output_hidden_states=True,)[-1] split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] return split_hidden_states def encode_prompt_edit(self, pipe: QwenImagePipeline, prompt, edit_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<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" drop_idx = 64 txt = [template.format(e) for e in prompt] model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] return split_hidden_states def encode_prompt_edit_multi(self, pipe: QwenImagePipeline, prompt, edit_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" drop_idx = 64 img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" base_img_prompt = "".join([img_prompt_template.format(i + 1) for i in range(len(edit_image))]) txt = [template.format(base_img_prompt + e) for e in prompt] edit_image = [self.resize_image(image) for image in edit_image] model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] return split_hidden_states def process(self, pipe: QwenImagePipeline, prompt, edit_image=None) -> dict: pipe.load_models_to_device(self.onload_model_names) if pipe.text_encoder is not None: prompt = [prompt] if edit_image is None: split_hidden_states = self.encode_prompt(pipe, prompt) elif isinstance(edit_image, Image.Image): split_hidden_states = self.encode_prompt_edit(pipe, prompt, edit_image) else: split_hidden_states = self.encode_prompt_edit_multi(pipe, prompt, edit_image) attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max([e.size(0) for e in split_hidden_states]) prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) prompt_embeds = prompt_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) return {"prompt_emb": prompt_embeds, "prompt_emb_mask": encoder_attention_mask} else: return {} class QwenImageUnit_EditImageEmbedder(PipelineUnit): def __init__(self): super().__init__( input_params=("edit_image", "tiled", "tile_size", "tile_stride", "edit_image_auto_resize"), output_params=("edit_latents", "edit_image"), onload_model_names=("vae",) ) def calculate_dimensions(self, target_area, ratio): import math width = math.sqrt(target_area * ratio) height = width / ratio width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height def edit_image_auto_resize(self, edit_image): calculated_width, calculated_height = self.calculate_dimensions(1024 * 1024, edit_image.size[0] / edit_image.size[1]) return edit_image.resize((calculated_width, calculated_height)) def process(self, pipe: QwenImagePipeline, edit_image, tiled, tile_size, tile_stride, edit_image_auto_resize=False): if edit_image is None: return {} if pipe.extractor is not None: return {} pipe.load_models_to_device(self.onload_model_names) if isinstance(edit_image, Image.Image): resized_edit_image = self.edit_image_auto_resize(edit_image) if edit_image_auto_resize else edit_image edit_image = pipe.preprocess_image(resized_edit_image).to(device=pipe.device, dtype=pipe.torch_dtype) edit_latents = pipe.vae.encode(edit_image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) else: resized_edit_image, edit_latents = [], [] for image in edit_image: if edit_image_auto_resize: image = self.edit_image_auto_resize(image) resized_edit_image.append(image) image=resized_edit_image[0] image = pipe.preprocess_image(image).to(device=pipe.device, dtype=pipe.torch_dtype) latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) edit_latents.append(latents) return {"edit_latents": edit_latents, "edit_image": resized_edit_image} class QwenImageUnit_MirrorPPROperationEncoder(PipelineUnit): def __init__(self): super().__init__( input_params=("example_origin","example_target"), output_params=("operation_emb"), onload_model_names=("mae","rformer","connector") ) def process_image_for_tokenizer(self,image_input, target_size=(1024, 1024)): """ Args: image_input: 单个 PIL.Image 对象，或者 PIL.Image 对象的列表。 target_size: 目标尺寸 (height, width)，默认 (1024, 1024)。 Returns: torch.Tensor: - 如果输入是单个图片: 返回 (3, H, W) - 如果输入是列表: 返回 (B, 3, H, W) - 数据类型 float32，已归一化。 """ mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) transforms = v2.Compose([ v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Resize(target_size, interpolation=InterpolationMode.BICUBIC, antialias=True), ]) def _process_one_tensor(img): x = transforms(img) if mean.device != x.device: mean_dev = mean.to(x.device) std_dev = std.to(x.device) else: mean_dev, std_dev = mean, std x = (x - mean_dev) / (std_dev + 1e-6) return x if isinstance(image_input, list): tensors = [_process_one_tensor(img) for img in image_input] return torch.stack(tensors) else: return _process_one_tensor(image_input).unsqueeze(0) def process(self, pipe: QwenImagePipeline,example_origin,example_target): if pipe.rformer2d is not None: return {} if pipe.extractor is not None: return {} pipe.load_models_to_device(self.onload_model_names) example_origin_pixel_values = self.process_image_for_tokenizer(example_origin).to(device=pipe.device, dtype=pipe.torch_dtype) example_target_pixel_values = self.process_image_for_tokenizer(example_target).to(device=pipe.device, dtype=pipe.torch_dtype) example_origin_hidden_states=pipe.mae(example_origin_pixel_values) example_target_hidden_states=pipe.mae(example_target_pixel_values) query_num = pipe.rformer.query_num latent_motion_tokens = pipe.rformer( cond_hidden_states=example_origin_hidden_states, target_hidden_states=example_target_hidden_states).last_hidden_state[:, :query_num] for layer in pipe.connector.encoder.layers: latent_motion_tokens=gradient_checkpoint_forward(layer,True,False,latent_motion_tokens) operation_emb=pipe.connector.proj(latent_motion_tokens) return {"operation_emb": operation_emb} class QwenImageUnit_MirrorPPROperationEncoder2D(PipelineUnit): def __init__(self): super().__init__( input_params=("example_origin","example_target","shape_1"), output_params=("operation_emb"), onload_model_names=("mae","rformer2d","connector") ) def process_image_for_tokenizer(self,image_input): """ Args: image_input: 单个 PIL.Image 对象，或者 PIL.Image 对象的列表。 Returns: torch.Tensor: - 如果输入是单个图片: 返回 (3, H, W) - 如果输入是列表: 返回 (B, 3, H, W) - 数据类型 float32，已归一化。 """ mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) transforms = v2.Compose([ v2.ToImage(), v2.ToDtype(torch.float32, scale=True), ]) def _process_one_tensor(img): x = transforms(img) if mean.device != x.device: mean_dev = mean.to(x.device) std_dev = std.to(x.device) else: mean_dev, std_dev = mean, std x = (x - mean_dev) / (std_dev + 1e-6) return x if isinstance(image_input, list): tensors = [_process_one_tensor(img) for img in image_input] return torch.stack(tensors) else: return _process_one_tensor(image_input).unsqueeze(0) def process(self, pipe: QwenImagePipeline,example_origin,example_target,shape_1): if pipe.rformer is not None: return {} if pipe.extractor is not None: return {} h1,w1=shape_1 new_shape_1 = [torch.tensor([h1]), torch.tensor([w1])] pipe.load_models_to_device(self.onload_model_names) example_origin_pixel_values = self.process_image_for_tokenizer(example_origin).to(device=pipe.device, dtype=pipe.torch_dtype) example_target_pixel_values = self.process_image_for_tokenizer(example_target).to(device=pipe.device, dtype=pipe.torch_dtype) example_origin_hidden_states=pipe.mae(example_origin_pixel_values) example_target_hidden_states=pipe.mae(example_target_pixel_values) query_num = pipe.rformer2d.query_num latent_motion_tokens = pipe.rformer2d( cond_hidden_states=example_origin_hidden_states, target_hidden_states=example_target_hidden_states,sample1_shapes=new_shape_1).last_hidden_state[:, :query_num] for layer in pipe.connector.encoder.layers: latent_motion_tokens=gradient_checkpoint_forward(layer,True,False,latent_motion_tokens) operation_emb=pipe.connector.proj(latent_motion_tokens) return {"operation_emb": operation_emb} class QwenImageUnit_PromptEmbedderWithExtractor(PipelineUnit): def __init__(self): super().__init__( input_params=("prompt", "edit_image", "tiled", "tile_size", "tile_stride", "edit_image_auto_resize"), output_params=("prompt_emb", "prompt_emb_mask", "image_vl_emb","edit_latents","edit_image"), onload_model_names=("vae", "extractor", "text_encoder") ) def extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() valid_lengths = bool_mask.sum(dim=1) selected = hidden_states[bool_mask] split_result = torch.split(selected, valid_lengths.tolist(), dim=0) return split_result def calculate_dimensions(self, target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height def edit_image_auto_resize(self, edit_image): calculated_width, calculated_height = self.calculate_dimensions(1024 * 1024, edit_image.size[0] / edit_image.size[1]) return edit_image.resize((calculated_width, calculated_height)) def encode_image_vl_emb(self, pipe: QwenImagePipeline, prompt, edit_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<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" drop_idx = 64 txt = [template.format(e) for e in prompt] model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) hidden_states = pipe.text_encoder(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw, output_hidden_states=True,)[-1] split_hidden_states = self.extract_masked_hidden(hidden_states, model_inputs.attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] return split_hidden_states def encode_prompt_edit(self, pipe: QwenImagePipeline, prompt, edit_image,extractor_emb): 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<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" drop_idx = 64 txt = [template.format(e) for e in prompt] model_inputs = pipe.processor(text=txt, images=edit_image, padding=True, return_tensors="pt").to(pipe.device) inputs_embeds,position_ids=pipe.text_encoder.vision_forward(input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw) batch_size = inputs_embeds.shape[0] num_queries = extractor_emb.shape[1] inputs_embeds = torch.cat([inputs_embeds, extractor_emb], dim=1) max_pos = position_ids.max(dim=-1, keepdim=True).values query_pos = torch.arange(1, num_queries + 1, device=position_ids.device).view(1, 1, -1).expand(3, batch_size, -1) position_ids = torch.cat([position_ids, max_pos + query_pos], dim=-1) attention_mask = torch.cat([model_inputs.attention_mask, torch.ones((batch_size, num_queries), device=model_inputs.attention_mask.device, dtype=model_inputs.attention_mask.dtype)], dim=1) hidden_states=pipe.text_encoder.llm_forward(inputs_embeds=inputs_embeds,position_ids=position_ids,attention_mask=attention_mask,output_hidden_states=True)[-1] return [h[-num_queries:] for h in hidden_states] def process(self, pipe: QwenImagePipeline, prompt, edit_image, tiled, tile_size, tile_stride, edit_image_auto_resize=False): if pipe.extractor is None: return {} print("using QwenImageUnit_PromptEmbedderWithExtractor") pipe.load_models_to_device(self.onload_model_names) if isinstance(edit_image, list): edit_images_list = edit_image else: edit_images_list = [edit_image] latents_list = [] processed_images = [] for img in edit_images_list: if edit_image_auto_resize: img = self.edit_image_auto_resize(img) processed_images.append(img) img_tensor = pipe.preprocess_image(img).to(device=pipe.device, dtype=pipe.torch_dtype) latent = pipe.vae.encode(img_tensor, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) latents_list.append(latent) latents_stacked = torch.stack(latents_list, dim=1) extractor_emb = pipe.extractor(latents_stacked) prompt1=[""] split_hidden_states = self.encode_prompt_edit(pipe, prompt1, edit_image[0],extractor_emb) prompt2=["Generate the retouched version of the input image."] image_split_hidden_states=self.encode_image_vl_emb(pipe, prompt2, edit_image[0]) attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max([e.size(0) for e in split_hidden_states]) prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) prompt_embeds = prompt_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) image_attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in image_split_hidden_states] image_max_seq_len = max([e.size(0) for e in image_split_hidden_states]) image_embeds = torch.stack([torch.cat([u, u.new_zeros(image_max_seq_len - u.size(0), u.size(1))]) for u in image_split_hidden_states]) image_encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(image_max_seq_len - u.size(0))]) for u in image_attn_mask_list]) image_embeds = image_embeds.to(dtype=pipe.torch_dtype, device=pipe.device) combined_attention_mask = torch.cat([image_encoder_attention_mask,encoder_attention_mask], dim=1) edit_latents=[latents_list[0]] return {"prompt_emb": prompt_embeds, "prompt_emb_mask": combined_attention_mask,"image_vl_emb":image_embeds,"edit_latents": edit_latents, "edit_image": processed_images} def model_fn_qwen_image( dit: QwenImageDiT = None, blockwise_controlnet: QwenImageBlockwiseMultiControlNet = None, latents=None, timestep=None, prompt_emb=None, prompt_emb_mask=None, operation_emb=None, height=None, width=None, blockwise_controlnet_conditioning=None, blockwise_controlnet_inputs=None, progress_id=0, num_inference_steps=1, entity_prompt_emb=None, entity_prompt_emb_mask=None, entity_masks=None, edit_latents=None, context_latents=None, enable_fp8_attention=False, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, edit_rope_interpolation=False, zero_cond_t=False, output_connector_embeds=False, no_qwenvl=False, **kwargs ): img_shapes = [(latents.shape[0], latents.shape[2]//2, latents.shape[3]//2)] timestep = timestep / 1000 image = rearrange(latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=height//16, W=width//16, P=2, Q=2) image_seq_len = image.shape[1] if context_latents is not None: img_shapes += [(context_latents.shape[0], context_latents.shape[2]//2, context_latents.shape[3]//2)] context_image = rearrange(context_latents, "B C (H P) (W Q) -> B (H W) (C P Q)", H=context_latents.shape[2]//2, W=context_latents.shape[3]//2, P=2, Q=2) image = torch.cat([image, context_image], dim=1) if edit_latents is not None: edit_latents_list = edit_latents if isinstance(edit_latents, list) else [edit_latents] img_shapes += [(e.shape[0], e.shape[2]//2, e.shape[3]//2) for e in edit_latents_list] edit_image = [rearrange(e, "B C (H P) (W Q) -> B (H W) (C P Q)", H=e.shape[2]//2, W=e.shape[3]//2, P=2, Q=2) for e in edit_latents_list] image = torch.cat([image] + edit_image, dim=1) image = dit.img_in(image) if zero_cond_t: timestep = torch.cat([timestep, timestep * 0], dim=0) modulate_index = torch.tensor( [[0] * prod(sample[0]) + [1] * sum([prod(s) for s in sample[1:]]) for sample in [img_shapes]], device=timestep.device, dtype=torch.int, ) else: modulate_index = None conditioning = dit.time_text_embed(timestep, image.dtype) if prompt_emb is not None and prompt_emb_mask is not None and not no_qwenvl: final_prompt_emb = torch.cat([prompt_emb,operation_emb], dim=1) txt_seq_lens = prompt_emb_mask.sum(dim=1).tolist() txt_seq_lens = [x + operation_emb.shape[1] for x in txt_seq_lens] else: final_prompt_emb = operation_emb txt_seq_lens = [operation_emb.shape[1]] * operation_emb.shape[0] if entity_prompt_emb is not None: text, image_rotary_emb, attention_mask = dit.process_entity_masks( latents, final_prompt_emb, prompt_emb_mask, entity_prompt_emb, entity_prompt_emb_mask, entity_masks, height, width, image, img_shapes, ) else: text = dit.txt_in(dit.txt_norm(final_prompt_emb)) if edit_rope_interpolation: image_rotary_emb = dit.pos_embed.forward_sampling(img_shapes, txt_seq_lens, device=latents.device) else: image_rotary_emb = dit.pos_embed(img_shapes, txt_seq_lens, device=latents.device) attention_mask = None if blockwise_controlnet_conditioning is not None: blockwise_controlnet_conditioning = blockwise_controlnet.preprocess( blockwise_controlnet_inputs, blockwise_controlnet_conditioning) for block_id, block in enumerate(dit.transformer_blocks): text, image = gradient_checkpoint_forward( block, use_gradient_checkpointing, use_gradient_checkpointing_offload, image=image, text=text, temb=conditioning, image_rotary_emb=image_rotary_emb, attention_mask=attention_mask, enable_fp8_attention=enable_fp8_attention, modulate_index=modulate_index, ) if blockwise_controlnet_conditioning is not None: image_slice = image[:, :image_seq_len].clone() controlnet_output = blockwise_controlnet.blockwise_forward( image=image_slice, conditionings=blockwise_controlnet_conditioning, controlnet_inputs=blockwise_controlnet_inputs, block_id=block_id, progress_id=progress_id, num_inference_steps=num_inference_steps, ) image[:, :image_seq_len] = image_slice + controlnet_output if zero_cond_t: conditioning = conditioning.chunk(2, dim=0)[0] image = dit.norm_out(image, conditioning) image = dit.proj_out(image) image = image[:, :image_seq_len] latents = rearrange(image, "B (H W) (C P Q) -> B C (H P) (W Q)", H=height//16, W=width//16, P=2, Q=2) if output_connector_embeds: return latents,prompt_emb else: return latents