import argparse import torch from llmga.llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN from llmga.llava.conversation import conv_templates, SeparatorStyle from llmga.llava.model.builder import load_pretrained_model from llmga.llava.utils import disable_torch_init from llmga.llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria import requests from PIL import Image from io import BytesIO import copy from llmga.diffusers.pipeline_semantic_stable_diffusion_img2img_solver_lpw_mask import SemanticStableDiffusionImg2ImgPipeline_DPMSolver from diffusers import StableDiffusionPipeline, AutoencoderKL from diffusers.schedulers import DDIMScheduler from llmga.diffusers.scheduling_dpmsolver_multistep_inject import DPMSolverMultistepSchedulerInject import random import cv2 import PIL from PIL import Image import os import numpy as np import json from tqdm import tqdm def read_json(file_path): with open(file_path, 'r', encoding='utf-8') as file: data = json.load(file) return data def write_json(file_path, data): with open(file_path, 'w', encoding='utf-8') as file: json.dump(data, file, ensure_ascii=False, indent=4) def load_image(image_file): if image_file.startswith('http') or image_file.startswith('https'): response = requests.get(image_file) image = Image.open(BytesIO(response.content)).convert('RGB') else: image = Image.open(image_file).convert('RGB') return image def randomize_seed_fn(seed, is_random): if is_random: seed = random.randint(0, np.iinfo(np.int32).max) return seed def seed_everything(seed): torch.manual_seed(seed) torch.cuda.manual_seed(seed) random.seed(seed) np.random.seed(seed) def crop_image(image): h, w, c = image.shape if h < w: offset = (w - h) // 2 image = image[:, offset:offset + h] elif w < h: offset = (h - w) // 2 image = image[offset:offset + w] image = np.array(Image.fromarray(image).resize((512, 512))) return image def main(args): # Model disable_torch_init() model_name = get_model_name_from_path(args.model_path) #print(f"Model name: {model_name}") tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device) vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16) pipe = SemanticStableDiffusionImg2ImgPipeline_DPMSolver.from_pretrained(args.sd_model_id,vae=vae,torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False).to("cuda") pipe.scheduler = DPMSolverMultistepSchedulerInject.from_pretrained(args.sd_model_id, subfolder="scheduler" , algorithm_type="sde-dpmsolver++", solver_order=2) def sample(zs, wts, mask_image, attention_store, text_cross_attention_maps, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1): latents = wts[-1].expand(1, -1, -1, -1) img, attention_store, text_cross_attention_maps = pipe( prompt=prompt_tar, init_latents=latents, guidance_scale=cfg_scale_tar, mask_image=mask_image, attention_store = attention_store, text_cross_attention_maps=text_cross_attention_maps, zs=zs, ) return img.images[0], attention_store, text_cross_attention_maps if 'llama-2' in model_name.lower(): conv_mode = "llava_llama_2" elif 'llama3' in model_name.lower(): conv_mode = "llama_3" elif "gemma" in model_name.lower(): conv_mode = "gemma" elif "qwen2" in model_name.lower(): conv_mode = "qwen_2" elif "phi-3" in model_name.lower(): conv_mode = "phi_3" elif "mistral" in model_name.lower(): conv_mode = "mistral_instruct" elif "v1.6-34b" in model_name.lower(): conv_mode = "chatml_direct" elif "v1" in model_name.lower(): conv_mode = "llava_v1" elif "mpt" in model_name.lower(): conv_mode = "mpt" else: conv_mode = "llava_v0" if args.conv_mode is not None and conv_mode != args.conv_mode: print('[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}'.format(conv_mode, args.conv_mode, args.conv_mode)) else: args.conv_mode = conv_mode image = load_image(args.image_file) #/Users/baixuehai/Downloads/2025_AAAI/test-llmga-sd15-editing.sh image=np.array(image) image_np=crop_image(image) image=Image.fromarray(image_np) #print(image_processor) image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().cuda() image_size = image.size data_path = '/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/magicbrush_dataset/genp2_4_single.json' save_image = '/home/zbz5349/WorkSpace/aigeeks/LLMGA/single_result' os.makedirs(save_image,exist_ok=True) # 若有x个指令那么生成x(single) + x(mix) + 1(all)张图片 data = read_json(data_path) for id in tqdm(range(2000)): conv = conv_templates[args.conv_mode].copy() if "mpt" in model_name.lower(): roles = ('user', 'assistant') else: roles = conv.roles img_path = data[id]["content"][0]["image"] image = load_image(img_path) image=np.array(image) image_np=crop_image(image) image=Image.fromarray(image_np) #print(image_processor) image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().cuda() image_size = image.size p = data[id]["content"][1]["text"] g = img_path g = g.split('/') s_img = f"{p[0]}_{g[-1]}" for i in range(2): if i==0: inp="Generate a similar image" else: inp = p[2:] ##input(f"{roles[0]}: ") #print(f"{roles[1]}: ", end="") if image is not None: # first message if model.config.mm_use_im_start_end: inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp else: inp = DEFAULT_IMAGE_TOKEN + '\n' + inp conv.append_message(conv.roles[0], inp) image = None else: # later messages conv.append_message(conv.roles[0], inp) conv.append_message(conv.roles[1], None) prompt = conv.get_prompt() input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda() stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2 if conv_mode == "gemma": stop_str = conv.sep2 keywords = [stop_str] stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) print(image_tensor.shape) with torch.inference_mode(): output_ids = model.generate( input_ids, images=image_tensor, image_sizes=[image_size], do_sample=True if args.temperature > 0 else False, temperature=args.temperature, max_new_tokens=args.max_new_tokens, pad_token_id=tokenizer.eos_token_id, use_cache=True) outputs = tokenizer.decode(output_ids[0]).strip() conv.messages[-1][-1] = outputs if conv_mode == "gemma": outputs=copy.deepcopy(outputs)[:-19] elif conv_mode == "llama_3": outputs=copy.deepcopy(outputs)[:-10] elif conv_mode == "llava_v1": outputs=copy.deepcopy(outputs)[:-4] elif conv_mode == "phi_3": outputs=copy.deepcopy(outputs)[:-7] elif conv_mode == "qwen_2": outputs=copy.deepcopy(outputs)[:-10] elif conv_mode == "mistral_instruct": outputs=copy.deepcopy(outputs)[:-5] else: outputs=copy.deepcopy(outputs) caption=copy.deepcopy(outputs) id1=caption.find("") num_space=12 id2=caption.find("") if id1==-1 and id2==-1: caption = caption elif id1==-1 and id2!=-1: caption = caption[:id2] elif id1!=-1 and id2==-1: caption = caption[id1+num_space:] else: caption = caption[id1+num_space:id2] if id1==-1: outputs=caption # print(caption) # else: # print(outputs) if i==0: src_prompt = caption else: tar_prompt = caption init_image = Image.fromarray(image_np) mask_image = pipe.generate_mask(image=init_image, source_prompt=src_prompt, target_prompt=tar_prompt,mask_thresholding_ratio=3.0) cv2.imwrite(os.path.join(args.save_path,"mask.png"),mask_image[0]*255) zs_tensor, wts_tensor = pipe.invert( image_path = image_np, source_prompt =src_prompt, source_guidance_scale= args.src_cfg_scale, num_inversion_steps = args.steps, skip = args.skip, eta = 1.0, ) wts = wts_tensor zs = zs_tensor pure_ddpm_img, attention_store, text_cross_attention_maps = sample(zs, wts, mask_image, attention_store=None, text_cross_attention_maps=None, prompt_tar=tar_prompt, skip=args.skip, cfg_scale_tar=args.tar_cfg_scale) #Image.fromarray(image_np).save(os.path.join(args.save_path,"input_image.png")) pure_ddpm_img.save(os.path.join(save_image,s_img)) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model-path", type=str, default="facebook/opt-350m") parser.add_argument("--model-base", type=str, default=None) parser.add_argument("--image-file", type=str, required=True) parser.add_argument("--save_path", type=str, required=True) parser.add_argument("--num-gpus", type=int, default=1) parser.add_argument("--conv-mode", type=str, default=None) parser.add_argument("--temperature", type=float, default=0.2) parser.add_argument("--max-new-tokens", type=int, default=512) parser.add_argument("--load-8bit", action="store_true") parser.add_argument("--load-4bit", action="store_true") parser.add_argument("--debug", action="store_true") parser.add_argument("--device", type=str, default="cuda") #**************SD Augs***************** parser.add_argument("--sd_model_id", type=str, required=True) parser.add_argument("--src_cfg_scale", type=float, default=3.5) parser.add_argument("--steps", type=int, default=50) parser.add_argument("--skip", type=int, default=25) parser.add_argument("--tar_cfg_scale", type=float, default=7.5) args = parser.parse_args() main(args)