Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| # Copyright (c) XiMing Xing. All rights reserved. | |
| # Author: XiMing Xing | |
| # Description: | |
| from PIL import Image | |
| from typing import Union, AnyStr, List | |
| from omegaconf.listconfig import ListConfig | |
| import diffusers | |
| import numpy as np | |
| from tqdm.auto import tqdm | |
| import torch | |
| from torchvision import transforms | |
| import clip | |
| from pytorch_svgrender.libs.engine import ModelState | |
| from pytorch_svgrender.painter.vectorfusion import LSDSPipeline, LSDSSDXLPipeline, Painter, PainterOptimizer | |
| from pytorch_svgrender.painter.vectorfusion import channel_saturation_penalty_loss as pixel_penalty_loss | |
| from pytorch_svgrender.painter.live import xing_loss_fn | |
| from pytorch_svgrender.plt import plot_img, plot_couple | |
| from pytorch_svgrender.token2attn.ptp_utils import view_images | |
| from pytorch_svgrender.diffusers_warp import init_StableDiffusion_pipeline, model2res | |
| class VectorFusionPipeline(ModelState): | |
| def __init__(self, args): | |
| assert args.x.style in ["iconography", "pixelart", "low-poly", "painting", "sketch", "ink"] | |
| logdir_ = f"sd{args.seed}-" \ | |
| f"{'scratch' if args.x.skip_live else 'baseline'}" \ | |
| f"-{args.x.model_id}" \ | |
| f"-{args.x.style}" \ | |
| f"-P{args.x.num_paths}" \ | |
| f"{'-RePath' if args.x.path_reinit.use else ''}" | |
| super().__init__(args, log_path_suffix=logdir_) | |
| # create log dir | |
| self.png_logs_dir = self.result_path / "png_logs" | |
| self.svg_logs_dir = self.result_path / "svg_logs" | |
| self.ft_png_logs_dir = self.result_path / "ft_png_logs" | |
| self.ft_svg_logs_dir = self.result_path / "ft_svg_logs" | |
| self.sd_sample_dir = self.result_path / 'sd_samples' | |
| self.reinit_dir = self.result_path / "reinit_logs" | |
| if self.accelerator.is_main_process: | |
| self.png_logs_dir.mkdir(parents=True, exist_ok=True) | |
| self.svg_logs_dir.mkdir(parents=True, exist_ok=True) | |
| self.ft_png_logs_dir.mkdir(parents=True, exist_ok=True) | |
| self.ft_svg_logs_dir.mkdir(parents=True, exist_ok=True) | |
| self.sd_sample_dir.mkdir(parents=True, exist_ok=True) | |
| self.reinit_dir.mkdir(parents=True, exist_ok=True) | |
| self.select_fpth = self.result_path / 'select_sample.png' | |
| # make video log | |
| self.make_video = self.args.mv | |
| if self.make_video: | |
| self.frame_idx = 0 | |
| self.frame_log_dir = self.result_path / "frame_logs" | |
| self.frame_log_dir.mkdir(parents=True, exist_ok=True) | |
| if self.x_cfg.model_id == "sdxl": | |
| # default LSDSSDXLPipeline scheduler is EulerDiscreteScheduler | |
| # when LSDSSDXLPipeline calls, scheduler.timesteps will change in step 4 | |
| # which causes problem in sds add_noise() function | |
| # because the random t may not in scheduler.timesteps | |
| custom_pipeline = LSDSSDXLPipeline | |
| custom_scheduler = diffusers.DPMSolverMultistepScheduler | |
| elif self.x_cfg.model_id == 'sd21': | |
| custom_pipeline = LSDSPipeline | |
| custom_scheduler = diffusers.DDIMScheduler | |
| else: # sd14, sd15 | |
| custom_pipeline = LSDSPipeline | |
| custom_scheduler = diffusers.PNDMScheduler | |
| self.diffusion = init_StableDiffusion_pipeline( | |
| self.x_cfg.model_id, | |
| custom_pipeline=custom_pipeline, | |
| custom_scheduler=custom_scheduler, | |
| device=self.device, | |
| local_files_only=not args.diffuser.download, | |
| force_download=args.diffuser.force_download, | |
| resume_download=args.diffuser.resume_download, | |
| ldm_speed_up=self.x_cfg.ldm_speed_up, | |
| enable_xformers=self.x_cfg.enable_xformers, | |
| gradient_checkpoint=self.x_cfg.gradient_checkpoint, | |
| lora_path=self.x_cfg.lora_path | |
| ) | |
| self.g_device = torch.Generator(device=self.device).manual_seed(args.seed) | |
| self.style = self.x_cfg.style | |
| if self.style in ["pixelart", "low-poly"]: | |
| self.x_cfg.path_schedule = 'list' | |
| self.x_cfg.schedule_each = list([args.x.grid]) | |
| if self.style == "pixelart": | |
| self.x_cfg.lr_stage_one.lr_schedule = False | |
| self.x_cfg.lr_stage_two.lr_schedule = False | |
| def get_path_schedule(self, schedule_each: Union[int, List]): | |
| if self.x_cfg.path_schedule == 'repeat': | |
| return int(self.x_cfg.num_paths / schedule_each) * [schedule_each] | |
| elif self.x_cfg.path_schedule == 'list': | |
| assert isinstance(self.x_cfg.schedule_each, list) or \ | |
| isinstance(self.x_cfg.schedule_each, ListConfig) | |
| return schedule_each | |
| else: | |
| raise NotImplementedError | |
| def target_file_preprocess(self, tar_path: AnyStr): | |
| process_comp = transforms.Compose([ | |
| transforms.Resize(size=(self.x_cfg.image_size, self.x_cfg.image_size)), | |
| transforms.ToTensor(), | |
| transforms.Lambda(lambda t: t.unsqueeze(0)), | |
| ]) | |
| tar_pil = Image.open(tar_path).convert("RGB") # open file | |
| target_img = process_comp(tar_pil) # preprocess | |
| target_img = target_img.to(self.device) | |
| return target_img | |
| def rejection_sampling(self, img_caption: Union[AnyStr, List], diffusion_samples: List): | |
| clip_model, preprocess = clip.load("ViT-B/32", device=self.device) | |
| text_input = clip.tokenize([img_caption]).to(self.device) | |
| text_features = clip_model.encode_text(text_input) | |
| text_features = text_features / text_features.norm(dim=-1, keepdim=True) | |
| clip_images = torch.stack([ | |
| preprocess(sample) for sample in diffusion_samples] | |
| ).to(self.device) | |
| image_features = clip_model.encode_image(clip_images) | |
| image_features = image_features / image_features.norm(dim=-1, keepdim=True) | |
| # clip score | |
| similarity_scores = (text_features @ image_features.T).squeeze(0) | |
| selected_image_index = similarity_scores.argmax().item() | |
| selected_image = diffusion_samples[selected_image_index] | |
| return selected_image | |
| def diffusion_sampling(self, text_prompt: AnyStr): | |
| """sampling K images""" | |
| diffusion_samples = [] | |
| for i in range(self.x_cfg.K): | |
| height = width = model2res(self.x_cfg.model_id) | |
| outputs = self.diffusion(prompt=[text_prompt], | |
| negative_prompt=self.args.neg_prompt, | |
| height=height, | |
| width=width, | |
| num_images_per_prompt=1, | |
| num_inference_steps=self.x_cfg.num_inference_steps, | |
| guidance_scale=self.x_cfg.guidance_scale, | |
| generator=self.g_device) | |
| outputs_np = [np.array(img) for img in outputs.images] | |
| view_images(outputs_np, save_image=True, fp=self.sd_sample_dir / f'samples_{i}.png') | |
| diffusion_samples.extend(outputs.images) | |
| self.print(f"num_generated_samples: {len(diffusion_samples)}, shape: {outputs_np[0].shape}") | |
| return diffusion_samples | |
| def LIVE_rendering(self, text_prompt: AnyStr): | |
| select_fpth = self.select_fpth | |
| # sampling K images | |
| diffusion_samples = self.diffusion_sampling(text_prompt) | |
| # rejection sampling | |
| select_target = self.rejection_sampling(text_prompt, diffusion_samples) | |
| select_target_pil = Image.fromarray(np.asarray(select_target)) # numpy to PIL | |
| select_target_pil.save(select_fpth) | |
| # load target file | |
| assert select_fpth.exists(), f"{select_fpth} is not exist!" | |
| target_img = self.target_file_preprocess(select_fpth.as_posix()) | |
| self.print(f"load target file from: {select_fpth.as_posix()}") | |
| # log path_schedule | |
| path_schedule = self.get_path_schedule(self.x_cfg.schedule_each) | |
| self.print(f"path_schedule: {path_schedule}") | |
| renderer = self.load_renderer() | |
| # first init center | |
| renderer.component_wise_path_init(target_img, pred=None, init_type=self.x_cfg.coord_init) | |
| optimizer_list = [PainterOptimizer(renderer, self.style, self.x_cfg.num_iter, | |
| self.x_cfg.lr_stage_one, self.x_cfg.trainable_bg) | |
| for _ in range(len(path_schedule))] | |
| pathn_record = [] | |
| loss_weight_keep = 0 | |
| total_step = len(path_schedule) * self.x_cfg.num_iter | |
| with tqdm(initial=self.step, total=total_step, disable=not self.accelerator.is_main_process) as pbar: | |
| for path_idx, pathn in enumerate(path_schedule): | |
| # record path | |
| pathn_record.append(pathn) | |
| # init graphic | |
| img = renderer.init_image(stage=0, num_paths=pathn) | |
| plot_img(img, self.result_path, fname=f"init_img_{path_idx}") | |
| # rebuild optimizer | |
| optimizer_list[path_idx].init_optimizers(pid_delta=int(path_idx * pathn)) | |
| pbar.write(f"=> adding {pathn} paths, n_path: {sum(pathn_record)}, " | |
| f"n_points: {len(renderer.get_point_parameters())}, " | |
| f"n_colors: {len(renderer.get_color_parameters())}") | |
| for t in range(self.x_cfg.num_iter): | |
| raster_img = renderer.get_image(step=t).to(self.device) | |
| if self.make_video and (self.step % self.args.framefreq == 0 or self.step == total_step - 1): | |
| plot_img(raster_img, self.frame_log_dir, fname=f"iter{self.frame_idx}") | |
| self.frame_idx += 1 | |
| if self.x_cfg.use_distance_weighted_loss and not (self.style == "pixelart"): | |
| loss_weight = renderer.calc_distance_weight(loss_weight_keep) | |
| # reconstruction loss | |
| if self.style == "pixelart": | |
| loss_recon = torch.nn.functional.l1_loss(raster_img, target_img) | |
| else: # UDF loss | |
| loss_recon = ((raster_img - target_img) ** 2) | |
| loss_recon = (loss_recon.sum(1) * loss_weight).mean() | |
| # Xing Loss for Self-Interaction Problem | |
| loss_xing = torch.tensor(0.) | |
| if self.style == "iconography": | |
| loss_xing = xing_loss_fn(renderer.get_point_parameters()) * self.x_cfg.xing_loss_weight | |
| # total loss | |
| loss = loss_recon + loss_xing | |
| lr_str = "" | |
| for k, lr in optimizer_list[path_idx].get_lr().items(): | |
| lr_str += f"{k}_lr: {lr:.4f}, " | |
| pbar.set_description( | |
| lr_str + | |
| f"L_total: {loss.item():.4f}, " | |
| f"L_recon: {loss_recon.item():.4f}, " | |
| f"L_xing: {loss_xing.item()}" | |
| ) | |
| # optimization | |
| for i in range(path_idx + 1): | |
| optimizer_list[i].zero_grad_() | |
| loss.backward() | |
| for i in range(path_idx + 1): | |
| optimizer_list[i].step_() | |
| renderer.clip_curve_shape() | |
| if self.x_cfg.lr_stage_one.lr_schedule: | |
| for i in range(path_idx + 1): | |
| optimizer_list[i].update_lr() | |
| if self.step % self.args.save_step == 0 and self.accelerator.is_main_process: | |
| plot_couple(target_img, | |
| raster_img, | |
| self.step, | |
| prompt=text_prompt, | |
| output_dir=self.png_logs_dir.as_posix(), | |
| fname=f"iter{self.step}") | |
| renderer.pretty_save_svg(self.svg_logs_dir / f"svg_iter{self.step}.svg") | |
| self.step += 1 | |
| pbar.update(1) | |
| # end a set of path optimization | |
| if self.x_cfg.use_distance_weighted_loss and not (self.style == "pixelart"): | |
| loss_weight_keep = loss_weight.detach().cpu().numpy() * 1 | |
| # recalculate the coordinates for the new join path | |
| renderer.component_wise_path_init(target_img, raster_img) | |
| # end LIVE | |
| final_svg_fpth = self.result_path / "live_stage_one_final.svg" | |
| renderer.pretty_save_svg(final_svg_fpth) | |
| if self.make_video: | |
| from subprocess import call | |
| call([ | |
| "ffmpeg", | |
| "-framerate", f"{self.args.framerate}", | |
| "-i", (self.frame_log_dir / "iter%d.png").as_posix(), | |
| "-vb", "20M", | |
| (self.result_path / "VF_rendering_stage1.mp4").as_posix() | |
| ]) | |
| return target_img, final_svg_fpth | |
| def painterly_rendering(self, text_prompt: AnyStr): | |
| # log prompts | |
| self.print(f"prompt: {text_prompt}") | |
| self.print(f"negative_prompt: {self.args.neg_prompt}\n") | |
| if self.x_cfg.skip_live: | |
| target_img = torch.randn(1, 3, self.x_cfg.image_size, self.x_cfg.image_size) | |
| final_svg_fpth = None | |
| self.print("from scratch with Score Distillation Sampling...") | |
| else: | |
| # text-to-img-to-svg | |
| target_img, final_svg_fpth = self.LIVE_rendering(text_prompt) | |
| torch.cuda.empty_cache() | |
| self.x_cfg.path_svg = final_svg_fpth | |
| self.print("\nfine-tune SVG via Score Distillation Sampling...") | |
| renderer = self.load_renderer(path_svg=final_svg_fpth) | |
| if self.x_cfg.skip_live: | |
| renderer.component_wise_path_init(target_img, pred=None, init_type='random') | |
| img = renderer.init_image(stage=0, num_paths=self.x_cfg.num_paths) | |
| plot_img(img, self.result_path, fname=f"init_img_stage_two") | |
| optimizer = PainterOptimizer(renderer, self.style, | |
| self.x_cfg.sds.num_iter, | |
| self.x_cfg.lr_stage_two, | |
| self.x_cfg.trainable_bg) | |
| optimizer.init_optimizers() | |
| self.print(f"-> Painter point Params: {len(renderer.get_point_parameters())}") | |
| self.print(f"-> Painter color Params: {len(renderer.get_color_parameters())}") | |
| self.print(f"-> Painter width Params: {len(renderer.get_width_parameters())}") | |
| self.step = 0 # reset global step | |
| total_step = self.x_cfg.sds.num_iter | |
| path_reinit = self.x_cfg.path_reinit | |
| self.print(f"\ntotal sds optimization steps: {total_step}") | |
| with tqdm(initial=self.step, total=total_step, disable=not self.accelerator.is_main_process) as pbar: | |
| while self.step < total_step: | |
| raster_img = renderer.get_image(step=self.step).to(self.device) | |
| if self.make_video and (self.step % self.args.framefreq == 0 or self.step == total_step - 1): | |
| plot_img(raster_img, self.frame_log_dir, fname=f"iter{self.frame_idx}") | |
| self.frame_idx += 1 | |
| L_sds, grad = self.diffusion.score_distillation_sampling( | |
| raster_img, | |
| im_size=self.x_cfg.sds.im_size, | |
| prompt=[text_prompt], | |
| negative_prompt=self.args.neg_prompt, | |
| guidance_scale=self.x_cfg.sds.guidance_scale, | |
| grad_scale=self.x_cfg.sds.grad_scale, | |
| t_range=list(self.x_cfg.sds.t_range), | |
| ) | |
| # Xing Loss for Self-Interaction Problem | |
| L_add = torch.tensor(0.) | |
| if self.style == "iconography": | |
| L_add = xing_loss_fn(renderer.get_point_parameters()) * self.x_cfg.xing_loss_weight | |
| # pixel_penalty_loss to combat oversaturation | |
| if self.style in ["pixelart", "low-poly"]: | |
| L_add = pixel_penalty_loss(raster_img) * self.x_cfg.penalty_weight | |
| loss = L_sds + L_add | |
| # optimization | |
| optimizer.zero_grad_() | |
| loss.backward() | |
| optimizer.step_() | |
| renderer.clip_curve_shape() | |
| # re-init paths | |
| if self.step % path_reinit.freq == 0 and self.step < path_reinit.stop_step and self.step != 0: | |
| renderer.reinitialize_paths(path_reinit.use, # on-off | |
| path_reinit.opacity_threshold, | |
| path_reinit.area_threshold, | |
| fpath=self.reinit_dir / f"reinit-{self.step}.svg") | |
| # update lr | |
| if self.x_cfg.lr_stage_two.lr_schedule: | |
| optimizer.update_lr() | |
| lr_str = "" | |
| for k, lr in optimizer.get_lr().items(): | |
| lr_str += f"{k}_lr: {lr:.4f}, " | |
| pbar.set_description( | |
| lr_str + | |
| f"L_total: {loss.item():.4f}, " | |
| f"L_add: {L_add.item():.4e}, " | |
| f"sds: {grad.item():.5e}" | |
| ) | |
| if self.step % self.args.save_step == 0 and self.accelerator.is_main_process: | |
| plot_couple(target_img, | |
| raster_img, | |
| self.step, | |
| prompt=text_prompt, | |
| output_dir=self.ft_png_logs_dir.as_posix(), | |
| fname=f"iter{self.step}") | |
| renderer.pretty_save_svg(self.ft_svg_logs_dir / f"svg_iter{self.step}.svg") | |
| self.step += 1 | |
| pbar.update(1) | |
| final_svg_fpth = self.result_path / "finetune_final.svg" | |
| renderer.pretty_save_svg(final_svg_fpth) | |
| if self.make_video: | |
| from subprocess import call | |
| call([ | |
| "ffmpeg", | |
| "-framerate", f"{self.args.framerate}", | |
| "-i", (self.frame_log_dir / "iter%d.png").as_posix(), | |
| "-vb", "20M", | |
| (self.result_path / "VF_rendering_stage2.mp4").as_posix() | |
| ]) | |
| self.close(msg="painterly rendering complete.") | |
| def load_renderer(self, path_svg=None): | |
| renderer = Painter(self.args.diffvg, | |
| self.style, | |
| self.x_cfg.num_segments, | |
| self.x_cfg.segment_init, | |
| self.x_cfg.radius, | |
| self.x_cfg.image_size, | |
| self.x_cfg.grid, | |
| self.x_cfg.trainable_bg, | |
| self.x_cfg.width, | |
| path_svg=path_svg, | |
| device=self.device) | |
| return renderer | |