pytorch_svgrender/pipelines/SVGDreamer_pipeline.py

# -*- coding: utf-8 -*-
# Copyright (c) XiMing Xing. All rights reserved.
# Author: XiMing Xing
# Description:
import pathlib
from PIL import Image
from typing import AnyStr

import numpy as np
from tqdm.auto import tqdm
import torch
from torch.optim.lr_scheduler import LambdaLR
import torchvision
from torchvision import transforms

from pytorch_svgrender.libs.engine import ModelState
from pytorch_svgrender.libs.solver.optim import get_optimizer
from pytorch_svgrender.painter.svgdreamer import Painter, PainterOptimizer
from pytorch_svgrender.painter.svgdreamer.painter_params import CosineWithWarmupLRLambda
from pytorch_svgrender.painter.live import xing_loss_fn
from pytorch_svgrender.painter.svgdreamer import VectorizedParticleSDSPipeline
from pytorch_svgrender.plt import plot_img
from pytorch_svgrender.utils.color_attrs import init_tensor_with_color
from pytorch_svgrender.token2attn.ptp_utils import view_images
from pytorch_svgrender.diffusers_warp import model2res

import ImageReward as RM


class SVGDreamerPipeline(ModelState):

    def __init__(self, args):
        assert args.x.style in ["iconography", "pixelart", "low-poly", "painting", "sketch", "ink"]
        assert args.x.guidance.n_particle >= args.x.guidance.vsd_n_particle
        assert args.x.guidance.n_particle >= args.x.guidance.phi_n_particle
        assert args.x.guidance.n_phi_sample >= 1

        logdir_ = f"sd{args.seed}" \
                  f"-{'vpsd' if args.x.skip_sive else 'sive'}" \
                  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"
        self.init_stage_two_dir = self.result_path / "stage_two_init_logs"
        self.phi_samples_dir = self.result_path / "phi_sampling_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.init_stage_two_dir.mkdir(parents=True, exist_ok=True)
            self.phi_samples_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)

        self.g_device = torch.Generator(device=self.device).manual_seed(args.seed)

        self.pipeline = VectorizedParticleSDSPipeline(args, args.diffuser, self.x_cfg.guidance, self.device)

        # load reward model
        self.reward_model = None
        if self.x_cfg.guidance.phi_ReFL:
            self.reward_model = RM.load("ImageReward-v1.0", device=self.device, download_root=self.x_cfg.reward_path)

        self.style = self.x_cfg.style
        if self.style == "pixelart":
            self.x_cfg.lr_stage_one.lr_schedule = False
            self.x_cfg.lr_stage_two.lr_schedule = False

    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 SIVE_stage(self, text_prompt: str):
        # TODO: SIVE implementation
        pass

    def painterly_rendering(self, text_prompt: str, target_file: AnyStr = None):
        # log prompts
        self.print(f"prompt: {text_prompt}")
        self.print(f"neg_prompt: {self.args.neg_prompt}\n")

        # for convenience
        im_size = self.x_cfg.image_size
        guidance_cfg = self.x_cfg.guidance
        n_particle = self.x_cfg.guidance.n_particle
        total_step = self.x_cfg.guidance.num_iter
        path_reinit = self.x_cfg.path_reinit

        init_from_target = True if (target_file and pathlib.Path(target_file).exists()) else False
        # switch mode
        if self.x_cfg.skip_sive and not init_from_target:
            # mode 1: optimization with VPSD from scratch
            # randomly init
            self.print("optimization with VPSD from scratch...")
            if self.x_cfg.color_init == 'rand':
                target_img = torch.randn(1, 3, im_size, im_size)
                self.print("color: randomly init")
            else:
                target_img = init_tensor_with_color(self.x_cfg.color_init, 1, im_size, im_size)
                self.print(f"color: {self.x_cfg.color_init}")

            # log init target_img
            plot_img(target_img, self.result_path, fname='init_target_img')
            final_svg_path = None
        elif init_from_target:
            # mode 2: load the SVG file and finetune it
            self.print(f"load svg from {target_file} ...")
            self.print(f"SVG fine-tuning via VPSD...")
            final_svg_path = target_file
            if self.x_cfg.color_init == 'target_randn':
                # special order: init newly paths color use random color
                target_img = torch.randn(1, 3, im_size, im_size)
                self.print("color: randomly init")
            else:
                # load the SVG and init newly paths color use target_img
                # note: the target will be converted to png via pydiffvg when load_renderer called
                target_img = None
        else:
            # mode 3: text-to-img-to-svg (two stage)
            target_img, final_svg_path = self.SIVE_stage(text_prompt)
            self.x_cfg.path_svg = final_svg_path
            self.print("\n SVG fine-tuning via VPSD...")
            plot_img(target_img, self.result_path, fname='init_target_img')

        # create svg renderer
        renderers = [self.load_renderer(final_svg_path) for _ in range(n_particle)]

        # randomly initialize the particles
        if self.x_cfg.skip_sive or init_from_target:
            if target_img is None:
                target_img = self.target_file_preprocess(self.result_path / 'target_img.png')
            for render in renderers:
                render.component_wise_path_init(gt=target_img, pred=None, init_type='random')

        # log init images
        for i, r in enumerate(renderers):
            init_imgs = r.init_image(stage=0, num_paths=self.x_cfg.num_paths)
            plot_img(init_imgs, self.init_stage_two_dir, fname=f"init_img_stage_two_{i}")

        # init renderer optimizer
        optimizers = []
        for renderer in renderers:
            optim_ = PainterOptimizer(renderer,
                                      self.style,
                                      guidance_cfg.num_iter,
                                      self.x_cfg.lr_stage_two,
                                      self.x_cfg.trainable_bg)
            optim_.init_optimizers()
            optimizers.append(optim_)

        # init phi_model optimizer
        phi_optimizer = get_optimizer('adamW',
                                      self.pipeline.phi_params,
                                      guidance_cfg.phi_lr,
                                      guidance_cfg.phi_optim)
        # init phi_model lr scheduler
        phi_scheduler = None
        schedule_cfg = guidance_cfg.phi_schedule
        if schedule_cfg.use:
            phi_lr_lambda = CosineWithWarmupLRLambda(num_steps=schedule_cfg.total_step,
                                                     warmup_steps=schedule_cfg.warmup_steps,
                                                     warmup_start_lr=schedule_cfg.warmup_start_lr,
                                                     warmup_end_lr=schedule_cfg.warmup_end_lr,
                                                     cosine_end_lr=schedule_cfg.cosine_end_lr)
            phi_scheduler = LambdaLR(phi_optimizer, lr_lambda=phi_lr_lambda, last_epoch=-1)

        self.print(f"-> Painter point Params: {len(renderers[0].get_point_parameters())}")
        self.print(f"-> Painter color Params: {len(renderers[0].get_color_parameters())}")
        self.print(f"-> Painter width Params: {len(renderers[0].get_width_parameters())}")

        L_reward = torch.tensor(0.)

        self.step = 0  # reset global step
        self.print(f"\ntotal VPSD 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:
                # set particles
                particles = [renderer.get_image() for renderer in renderers]
                raster_imgs = torch.cat(particles, dim=0)

                if self.make_video and (self.step % self.args.framefreq == 0 or self.step == total_step - 1):
                    plot_img(raster_imgs, self.frame_log_dir, fname=f"iter{self.frame_idx}")
                    self.frame_idx += 1

                L_guide, grad, latents, t_step = self.pipeline.variational_score_distillation(
                    raster_imgs,
                    self.step,
                    prompt=[text_prompt],
                    negative_prompt=self.args.neg_prompt,
                    grad_scale=guidance_cfg.grad_scale,
                    enhance_particle=guidance_cfg.particle_aug,
                    im_size=model2res(self.x_cfg.model_id)
                )

                # Xing Loss for Self-Interaction Problem
                L_add = torch.tensor(0.)
                if self.style == "iconography" or self.x_cfg.xing_loss.use:
                    for r in renderers:
                        L_add += xing_loss_fn(r.get_point_parameters()) * self.x_cfg.xing_loss.weight

                loss = L_guide + L_add

                # optimization
                for opt_ in optimizers:
                    opt_.zero_grad_()
                loss.backward()
                for opt_ in optimizers:
                    opt_.step_()

                # phi_model optimization
                for _ in range(guidance_cfg.phi_update_step):
                    L_lora = self.pipeline.train_phi_model(latents, guidance_cfg.phi_t, as_latent=True)

                    phi_optimizer.zero_grad()
                    L_lora.backward()
                    phi_optimizer.step()

                # reward learning
                if guidance_cfg.phi_ReFL and self.step % guidance_cfg.phi_sample_step == 0:
                    with torch.no_grad():
                        phi_outputs = []
                        phi_sample_paths = []
                        for idx in range(guidance_cfg.n_phi_sample):
                            phi_output = self.pipeline.sample(text_prompt,
                                                              num_inference_steps=guidance_cfg.phi_infer_step,
                                                              generator=self.g_device)
                            sample_path = (self.phi_samples_dir / f'iter{idx}.png').as_posix()
                            phi_output.images[0].save(sample_path)
                            phi_sample_paths.append(sample_path)

                            phi_output_np = np.array(phi_output.images[0])
                            phi_outputs.append(phi_output_np)
                        # save all samples
                        view_images(phi_outputs, save_image=True,
                                    num_rows=max(len(phi_outputs) // 6, 1),
                                    fp=self.phi_samples_dir / f'samples_iter{self.step}.png')

                    ranking, rewards = self.reward_model.inference_rank(text_prompt, phi_sample_paths)
                    self.print(f"ranking: {ranking}, reward score: {rewards}")

                    for k in range(guidance_cfg.n_phi_sample):
                        phi = self.target_file_preprocess(phi_sample_paths[ranking[k] - 1])
                        L_reward = self.pipeline.train_phi_model_refl(phi, weight=rewards[k])

                        phi_optimizer.zero_grad()
                        L_reward.backward()
                        phi_optimizer.step()

                # update the learning rate of the phi_model
                if phi_scheduler is not None:
                    phi_scheduler.step()

                # curve regularization
                for r in renderers:
                    r.clip_curve_shape()

                # re-init paths
                if self.step % path_reinit.freq == 0 and self.step < path_reinit.stop_step and self.step != 0:
                    for i, r in enumerate(renderers):
                        r.reinitialize_paths(path_reinit.use,  # on-off
                                             path_reinit.opacity_threshold,
                                             path_reinit.area_threshold,
                                             fpath=self.reinit_dir / f"reinit-{self.step}_p{i}.svg")

                # update lr
                if self.x_cfg.lr_stage_two.lr_schedule:
                    for opt_ in optimizers:
                        opt_.update_lr()

                # log pretrained model lr
                lr_str = ""
                for k, lr in optimizers[0].get_lr().items():
                    lr_str += f"{k}_lr: {lr:.4f}, "
                # log phi model lr
                cur_phi_lr = phi_optimizer.param_groups[0]['lr']
                lr_str += f"phi_lr: {cur_phi_lr:.3e}, "

                pbar.set_description(
                    lr_str +
                    f"t: {t_step.item():.2f}, "
                    f"L_total: {loss.item():.4f}, "
                    f"L_add: {L_add.item():.4e}, "
                    f"L_lora: {L_lora.item():.4f}, "
                    f"L_reward: {L_reward.item():.4f}, "
                    f"vpsd: {grad.item():.4e}"
                )

                if self.step % self.args.save_step == 0 and self.accelerator.is_main_process:
                    # save png
                    torchvision.utils.save_image(raster_imgs,
                                                 fp=self.ft_png_logs_dir / f'iter{self.step}.png')

                    # save svg
                    for i, r in enumerate(renderers):
                        r.pretty_save_svg(self.ft_svg_logs_dir / f"svg_iter{self.step}_p{i}.svg")

                self.step += 1
                pbar.update(1)

        # save final
        for i, r in enumerate(renderers):
            final_svg_path = self.result_path / f"finetune_final_p_{i}.svg"
            r.pretty_save_svg(final_svg_path)
        # save SVGs
        torchvision.utils.save_image(raster_imgs, fp=self.result_path / f'all_particles.png')

        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 / "svgdreamer_rendering.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)

        # if load a svg file, then rasterize it
        save_path = self.result_path / 'target_img.png'
        if path_svg is not None and (not save_path.exists()):
            canvas_width, canvas_height, shapes, shape_groups = renderer.load_svg(path_svg)
            render_img = renderer.render_image(canvas_width, canvas_height, shapes, shape_groups)
            torchvision.utils.save_image(render_img, fp=save_path)
        return renderer