gracekim0513/Step1X-Edit
0
1import argparse2import datetime3import json 4import itertools5import math6import os7import spaces8import time9from pathlib import Path10 11 12import gradio as gr13import numpy as np14import torch15from einops import rearrange, repeat16from huggingface_hub import snapshot_download17from PIL import Image, ImageOps18from safetensors.torch import load_file19from torchvision.transforms import functional as F20from tqdm import tqdm 21 22import sampling23from modules.autoencoder import AutoEncoder24from modules.conditioner import Qwen25VL_7b_Embedder as Qwen2VLEmbedder25from modules.model_edit import Step1XParams, Step1XEdit26from diffusers import StableDiffusionPipeline27 28pipe = StableDiffusionPipeline.from_pretrained(29 "stabilityai/stable-diffusion-2",30 cache_dir="/tmp", # 작게 제한된 디렉토리 사용31 local_files_only=False32)33 34print("TORCH_CUDA", torch.cuda.is_available())35 36examples = [37["examples 2/meme.jpg", "turn into an illustration in studio ghibli style",("examples 2/meme.jpg","examples 2/ghibli_meme.jpg"),],38["examples 2/celeb_meme.jpg", "replace the gray blazer with a leather jacket",("examples 2/celeb_meme.jpg","examples 2/leather.jpg")],39["examples 2/cookie.png", "remove the cookie",("examples 2/cookie.png","examples 2/no_cookie.png")],40["examples 2/poster_orig.jpg", "replace 'lambs' with 'llamas'",("examples 2/poster_orig.jpg","examples 2/poster.jpg")],41]42 43def generate_examples(init_image, prompt):44 return inference(prompt, init_image, seed=-1, size_level=512)45 46 47def load_state_dict(model, ckpt_path, device="cuda", strict=False, assign=True):48 if Path(ckpt_path).suffix == ".safetensors":49 state_dict = load_file(ckpt_path, device)50 else:51 state_dict = torch.load(ckpt_path, map_location="cpu")52 53 missing, unexpected = model.load_state_dict(54 state_dict, strict=strict, assign=assign55 )56 if len(missing) > 0 and len(unexpected) > 0:57 print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))58 print("\n" + "-" * 79 + "\n")59 print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))60 elif len(missing) > 0:61 print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))62 elif len(unexpected) > 0:63 print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))64 return model65 66 67def load_models(68 dit_path=None,69 ae_path=None,70 qwen2vl_model_path=None,71 device="cuda",72 max_length=256,73 dtype=torch.bfloat16,74):75 qwen2vl_encoder = Qwen2VLEmbedder(76 qwen2vl_model_path,77 device=device,78 max_length=max_length,79 dtype=dtype,80 )81 82 with torch.device("meta"):83 ae = AutoEncoder(84 resolution=256,85 in_channels=3,86 ch=128,87 out_ch=3,88 ch_mult=[1, 2, 4, 4],89 num_res_blocks=2,90 z_channels=16,91 scale_factor=0.3611,92 shift_factor=0.1159,93 )94 95 step1x_params = Step1XParams(96 in_channels=64,97 out_channels=64,98 vec_in_dim=768,99 context_in_dim=4096,100 hidden_size=3072,101 mlp_ratio=4.0,102 num_heads=24,103 depth=19,104 depth_single_blocks=38,105 axes_dim=[16, 56, 56],106 theta=10_000,107 qkv_bias=True,108 )109 dit = Step1XEdit(step1x_params)110 111 ae = load_state_dict(ae, ae_path)112 dit = load_state_dict(113 dit, dit_path114 )115 116 dit = dit.to(device=device, dtype=dtype)117 ae = ae.to(device=device, dtype=torch.float32)118 119 return ae, dit, qwen2vl_encoder120 121 122class ImageGenerator:123 def __init__(124 self,125 dit_path=None,126 ae_path=None,127 qwen2vl_model_path=None,128 device="cuda",129 max_length=640,130 dtype=torch.bfloat16,131 ) -> None:132 self.device = torch.device(device)133 self.ae, self.dit, self.llm_encoder = load_models(134 dit_path=dit_path,135 ae_path=ae_path,136 qwen2vl_model_path=qwen2vl_model_path,137 max_length=max_length,138 dtype=dtype,139 )140 self.ae = self.ae.to(device=self.device, dtype=torch.float32)141 self.dit = self.dit.to(device=self.device, dtype=dtype)142 self.llm_encoder = self.llm_encoder.to(device=self.device, dtype=dtype)143 144 def to_cuda(self):145 self.ae.to(device='cuda', dtype=torch.float32)146 self.dit.to(device='cuda', dtype=torch.bfloat16)147 self.llm_encoder.to(device='cuda', dtype=torch.bfloat16)148 149 def prepare(self, prompt, img, ref_image, ref_image_raw):150 bs, _, h, w = img.shape151 bs, _, ref_h, ref_w = ref_image.shape152 153 assert h == ref_h and w == ref_w154 155 if bs == 1 and not isinstance(prompt, str):156 bs = len(prompt)157 elif bs >= 1 and isinstance(prompt, str):158 prompt = [prompt] * bs159 160 img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)161 ref_img = rearrange(ref_image, "b c (ref_h ph) (ref_w pw) -> b (ref_h ref_w) (c ph pw)", ph=2, pw=2)162 if img.shape[0] == 1 and bs > 1:163 img = repeat(img, "1 ... -> bs ...", bs=bs)164 ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)165 166 img_ids = torch.zeros(h // 2, w // 2, 3)167 168 img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]169 img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]170 img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)171 172 ref_img_ids = torch.zeros(ref_h // 2, ref_w // 2, 3)173 174 ref_img_ids[..., 1] = ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None]175 ref_img_ids[..., 2] = ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :]176 ref_img_ids = repeat(ref_img_ids, "ref_h ref_w c -> b (ref_h ref_w) c", b=bs)177 178 if isinstance(prompt, str):179 prompt = [prompt]180 181 txt, mask = self.llm_encoder(prompt, ref_image_raw)182 183 txt_ids = torch.zeros(bs, txt.shape[1], 3)184 185 img = torch.cat([img, ref_img.to(device=img.device, dtype=img.dtype)], dim=-2)186 img_ids = torch.cat([img_ids, ref_img_ids], dim=-2)187 188 189 return {190 "img": img,191 "mask": mask,192 "img_ids": img_ids.to(img.device),193 "llm_embedding": txt.to(img.device),194 "txt_ids": txt_ids.to(img.device),195 }196 197 @staticmethod198 def process_diff_norm(diff_norm, k):199 pow_result = torch.pow(diff_norm, k)200 201 result = torch.where(202 diff_norm > 1.0,203 pow_result,204 torch.where(diff_norm < 1.0, torch.ones_like(diff_norm), diff_norm),205 )206 return result207 208 def denoise(209 self,210 img: torch.Tensor,211 img_ids: torch.Tensor,212 llm_embedding: torch.Tensor,213 txt_ids: torch.Tensor,214 timesteps: list[float],215 cfg_guidance: float = 4.5,216 mask=None,217 show_progress=False,218 timesteps_truncate=1.0,219 ):220 if show_progress:221 pbar = tqdm(itertools.pairwise(timesteps), desc='denoising...')222 else:223 pbar = itertools.pairwise(timesteps)224 for t_curr, t_prev in pbar:225 if img.shape[0] == 1 and cfg_guidance != -1:226 img = torch.cat([img, img], dim=0)227 t_vec = torch.full(228 (img.shape[0],), t_curr, dtype=img.dtype, device=img.device229 )230 231 txt, vec = self.dit.connector(llm_embedding, t_vec, mask)232 233 234 pred = self.dit(235 img=img,236 img_ids=img_ids,237 txt=txt,238 txt_ids=txt_ids,239 y=vec,240 timesteps=t_vec,241 )242 243 if cfg_guidance != -1:244 cond, uncond = (245 pred[0 : pred.shape[0] // 2, :],246 pred[pred.shape[0] // 2 :, :],247 )248 if t_curr > timesteps_truncate:249 diff = cond - uncond250 diff_norm = torch.norm(diff, dim=(2), keepdim=True)251 pred = uncond + cfg_guidance * (252 cond - uncond253 ) / self.process_diff_norm(diff_norm, k=0.4)254 else:255 pred = uncond + cfg_guidance * (cond - uncond)256 tem_img = img[0 : img.shape[0] // 2, :] + (t_prev - t_curr) * pred257 img_input_length = img.shape[1] // 2258 img = torch.cat(259 [260 tem_img[:, :img_input_length],261 img[ : img.shape[0] // 2, img_input_length:],262 ], dim=1263 )264 265 return img[:, :img.shape[1] // 2]266 267 @staticmethod268 def unpack(x: torch.Tensor, height: int, width: int) -> torch.Tensor:269 return rearrange(270 x,271 "b (h w) (c ph pw) -> b c (h ph) (w pw)",272 h=math.ceil(height / 16),273 w=math.ceil(width / 16),274 ph=2,275 pw=2,276 )277 278 @staticmethod279 def load_image(image):280 from PIL import Image281 282 if isinstance(image, np.ndarray):283 image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0284 image = image.unsqueeze(0)285 return image286 elif isinstance(image, Image.Image):287 image = F.to_tensor(image.convert("RGB"))288 image = image.unsqueeze(0)289 return image290 elif isinstance(image, torch.Tensor):291 return image292 elif isinstance(image, str):293 image = F.to_tensor(Image.open(image).convert("RGB"))294 image = image.unsqueeze(0)295 return image296 else:297 raise ValueError(f"Unsupported image type: {type(image)}")298 299 def output_process_image(self, resize_img, image_size):300 res_image = resize_img.resize(image_size)301 return res_image302 303 def input_process_image(self, img, img_size=512):304 # 1. 打开图片305 w, h = img.size306 r = w / h 307 308 if w > h:309 w_new = math.ceil(math.sqrt(img_size * img_size * r))310 h_new = math.ceil(w_new / r)311 else:312 h_new = math.ceil(math.sqrt(img_size * img_size / r))313 w_new = math.ceil(h_new * r)314 h_new = math.ceil(h_new) // 16 * 16315 w_new = math.ceil(w_new) // 16 * 16316 317 img_resized = img.resize((w_new, h_new))318 return img_resized, img.size319 320 @torch.inference_mode()321 def generate_image(322 self,323 prompt,324 negative_prompt,325 ref_images,326 num_steps,327 cfg_guidance,328 seed,329 num_samples=1,330 init_image=None,331 image2image_strength=0.0,332 show_progress=False,333 size_level=512,334 ):335 assert num_samples == 1, "num_samples > 1 is not supported yet."336 ref_images_raw, img_info = self.input_process_image(ref_images, img_size=size_level)337 338 width, height = ref_images_raw.width, ref_images_raw.height339 340 341 ref_images_raw = self.load_image(ref_images_raw)342 ref_images_raw = ref_images_raw.to(self.device)343 # print(f'self.ae, self.dit device: {self.ae.device}, {self.dit.device}')344 ref_images = self.ae.encode(ref_images_raw.to(self.device) * 2 - 1)345 346 seed = int(seed)347 seed = torch.Generator(device="cpu").seed() if seed < 0 else seed348 349 t0 = time.perf_counter()350 351 if init_image is not None:352 init_image = self.load_image(init_image)353 init_image = init_image.to(self.device)354 init_image = torch.nn.functional.interpolate(init_image, (height, width))355 init_image = self.ae.encode(init_image.to() * 2 - 1)356 357 x = torch.randn(358 num_samples,359 16,360 height // 8,361 width // 8,362 device=self.device,363 dtype=torch.bfloat16,364 generator=torch.Generator(device=self.device).manual_seed(seed),365 )366 367 timesteps = sampling.get_schedule(368 num_steps, x.shape[-1] * x.shape[-2] // 4, shift=True369 )370 371 if init_image is not None:372 t_idx = int((1 - image2image_strength) * num_steps)373 t = timesteps[t_idx]374 timesteps = timesteps[t_idx:]375 x = t * x + (1.0 - t) * init_image.to(x.dtype)376 377 x = torch.cat([x, x], dim=0)378 ref_images = torch.cat([ref_images, ref_images], dim=0)379 ref_images_raw = torch.cat([ref_images_raw, ref_images_raw], dim=0)380 inputs = self.prepare([prompt, negative_prompt], x, ref_image=ref_images, ref_image_raw=ref_images_raw)381 382 x = self.denoise(383 **inputs,384 cfg_guidance=cfg_guidance,385 timesteps=timesteps,386 show_progress=show_progress,387 timesteps_truncate=1.0,388 )389 x = self.unpack(x.float(), height, width)390 with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):391 x = self.ae.decode(x)392 x = x.clamp(-1, 1)393 x = x.mul(0.5).add(0.5)394 395 t1 = time.perf_counter()396 print(f"Done in {t1 - t0:.1f}s.")397 images_list = []398 for img in x.float():399 images_list.append(self.output_process_image(F.to_pil_image(img), img_info))400 return images_list401 402 403# 模型仓库ID（如："bert-base-uncased"）404model_repo = "stepfun-ai/Step1X-Edit"405# 本地保存路径406model_path = "./model_weights"407os.makedirs(model_path, exist_ok=True)408 409 410# 下载模型（包括所有文件）411snapshot_download(412 repo_id=model_repo,413 local_dir=model_path,414 local_dir_use_symlinks=False # 避免使用符号链接415)416 417 418image_edit = ImageGenerator(419 ae_path=os.path.join(model_path, 'vae.safetensors'),420 dit_path=os.path.join(model_path, "step1x-edit-i1258.safetensors"),421 qwen2vl_model_path='Qwen/Qwen2.5-VL-7B-Instruct',422 max_length=640,423)424 425 426 427@spaces.GPU(duration=240)428def inference(prompt, ref_images, seed, size_level):429 start_time = time.time()430 431 if seed == -1:432 import random 433 random_seed = random.randint(0, 2**32 - 1)434 else:435 random_seed = seed436 437 image_edit.to_cuda()438 439 inference_func = image_edit.generate_image440 441 image = inference_func(442 prompt,443 negative_prompt="",444 ref_images=ref_images.convert('RGB'),445 num_samples=1,446 num_steps=28,447 cfg_guidance=6.0,448 seed=random_seed,449 show_progress=True,450 size_level=size_level,451 )[0]452 453 print(f"Time taken: {time.time() - start_time:.2f} seconds")454 return (ref_images, image), random_seed455 456with gr.Blocks() as demo:457 gr.Markdown(458 """459 # Step1X-Edit460 """461 )462 with gr.Row():463 with gr.Column():464 prompt = gr.Textbox(465 label="编辑指令 prompt",466 value='Remove the person from the image.',467 )468 init_image = gr.Image(label="Input Image", type='pil')469 470 random_seed = gr.Number(label="Random Seed", value=-1, minimum=-1)471 472 size_level = gr.Number(label="size level (recommend 512, 768, 1024, min 512)", value=512, minimum=512, maximum=1024)473 474 generate_btn = gr.Button("Generate")475 476 with gr.Column():477 output_image = gr.ImageSlider(label="Generated Image", type="pil", image_mode='RGB')478 output_random_seed = gr.Textbox(label="Used Seed", lines=5)479 from functools import partial480 generate_btn.click(481 fn=inference,482 inputs=[483 prompt, 484 init_image,485 random_seed,486 size_level,487 ],488 outputs=[output_image, output_random_seed],489 )490 491 gr.Examples(492 examples,493 inputs=[init_image, prompt],494 outputs=[output_image, output_random_seed],495 fn=generate_examples,496 cache_examples=True497 )498 499demo.launch()500