main2.py

import os
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
import cv2
import time
import tqdm
import numpy as np
import dearpygui.dearpygui as dpg

import torch
import torch.nn.functional as F

import trimesh
import rembg
import common.global_constant
from cam_utils import orbit_camera, OrbitCamera
from mesh_renderer import Renderer


# from kiui.lpips import LPIPS

class GUI:
    def __init__(self, opt):
        self.opt = opt  # shared with the trainer's opt to support in-place modification of rendering parameters.
        self.gui = opt.gui  # enable gui
        self.W = opt.W
        self.H = opt.H
        self.cam = OrbitCamera(opt.W, opt.H, r=opt.radius, fovy=opt.fovy)
        self.fixed_cam = None
        self.fixed_cam_list = []
        self.mode = "image"
        self.seed = "random"

        # only for gui
        self.buffer_image = np.ones((self.W, self.H, 3), dtype=np.float32)
        self.need_update = True  # update buffer_image

        # models
        self.device = torch.device(common.global_constant.GPU.DEVICE)
        self.bg_remover = None

        self.guidance_sd = None
        self.guidance_zero123 = None

        self.enable_sd = False
        self.enable_zero123 = False

        # renderer NN Module
        self.renderer = Renderer(opt).to(self.device)

        # input image
        self.input_img = None
        self.input_mask = None
        self.input_img_torch = None
        self.input_mask_torch = None
        # {{ only for GUI
        self.overlay_input_img = False
        self.overlay_input_img_ratio = 0.5
        # }}
        self.input_img_list = []
        self.input_mask_list = []
        self.input_img_torch_list = []
        self.input_mask_torch_list = []
        self.input_img_torch_channel_last_list = []

        # input text
        self.prompt = ""
        self.negative_prompt = ""
        self.prompt_list = []
        self.negative_prompt_list = []

        # training stuff
        self.training = False
        self.optimizer = None
        self.step = 0
        self.train_steps = 1  # steps per rendering loop
        # self.lpips_loss = LPIPS(net='vgg').to(self.device)

        # load input data from cmdline
        if self.opt.input is not None:
            self.load_input(self.opt.input)

        if self.opt.input_files is not None:
            input_files = self.opt.input_files.split(',')
            self.load_input_files(input_files)
        else:
            if self.input_img is not None:
                self.input_img_list.append(self.input_img)
                self.input_mask_list.append(self.input_mask)
        # print(self.input_img, self.input_mask)
        # print(self.input_img_list, self.input_mask_list, self.prompt_list)

        # override prompt from cmdline
        if self.opt.prompt is not None:
            self.prompt = self.opt.prompt
        if self.opt.negative_prompt is not None:
            self.negative_prompt = self.opt.negative_prompt
        print('prompt:', self.prompt, '; negative_prompt:', self.negative_prompt)
        # exit(0)

        if self.gui:
            dpg.create_context()
            self.register_dpg()
            self.test_step()

    def __del__(self):
        if self.gui:
            dpg.destroy_context()

    def seed_everything(self):
        try:
            seed = int(self.seed)
        except:
            seed = np.random.randint(0, 1000000)

        os.environ["PYTHONHASHSEED"] = str(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = True

        self.last_seed = seed

    def prepare_train(self):

        self.step = 0

        # setup training
        self.optimizer = torch.optim.Adam(self.renderer.get_params())

        # default camera
        if self.opt.mvdream or self.opt.imagedream:
            # the second view is the front view for mvdream/imagedream.
            pose = orbit_camera(self.opt.elevation, 90, self.opt.radius)
        else:
            pose = orbit_camera(self.opt.elevation, 0, self.opt.radius)
        self.fixed_cam = (pose, self.cam.perspective)
        # TODO: self.fixed_cam_list
        if self.opt.input_camera_pose is not None:
            input_camera_poses = self.opt.input_camera_pose.split(',')
            for pose in input_camera_poses:
                pose = float(pose)
                print('pose: %s' % pose)
                if self.opt.mvdream or self.opt.imagedream:
                    # the second view is the front view for mvdream/imagedream.
                    pose = orbit_camera(self.opt.elevation, 90 + pose, self.opt.radius)
                else:
                    pose = orbit_camera(self.opt.elevation, pose, self.opt.radius)
                self.fixed_cam_list.append((pose, self.cam.perspective))
        else:
            self.fixed_cam_list.append(self.fixed_cam)

        self.enable_sd = self.opt.lambda_sd > 0 and self.prompt != ""
        self.enable_zero123 = self.opt.lambda_zero123 > 0 and (
                self.input_img is not None or len(self.input_img_list) > 0)

        # lazy load guidance model
        if self.guidance_sd is None and self.enable_sd:
            if self.opt.mvdream:
                print(f"[INFO] loading MVDream...")
                from guidance.mvdream_utils import MVDream
                self.guidance_sd = MVDream(self.device)
                print(f"[INFO] loaded MVDream!")
            elif self.opt.imagedream:
                print(f"[INFO] loading ImageDream...")
                from guidance.imagedream_utils import ImageDream
                self.guidance_sd = ImageDream(self.device)
                print(f"[INFO] loaded ImageDream!")
            else:
                print(f"[INFO] loading SD...")
                from guidance.sd_utils import StableDiffusion
                self.guidance_sd = StableDiffusion(self.device)
                print(f"[INFO] loaded SD!")

        if self.guidance_zero123 is None and self.enable_zero123:
            print(f"[INFO] loading zero123...")
            model_key = None
            from guidance.zero123_utils import Zero123
            if self.opt.stable_zero123:
                model_key = 'ashawkey/stable-zero123-diffusers'
                self.guidance_zero123 = Zero123(self.device, model_key=model_key)
            else:
                model_key = 'ashawkey/zero123-xl-diffusers'
                self.guidance_zero123 = Zero123(self.device, model_key=model_key)
            print(f"[INFO] loaded zero123!", model_key)

        # input image
        '''
        if self.input_img is not None:
            self.input_img_torch = torch.from_numpy(self.input_img).permute(2, 0, 1).unsqueeze(0).to(self.device)
            self.input_img_torch = F.interpolate(self.input_img_torch, (self.opt.ref_size, self.opt.ref_size),
                                                 mode="bilinear", align_corners=False)

            self.input_mask_torch = torch.from_numpy(self.input_mask).permute(2, 0, 1).unsqueeze(0).to(self.device)
            self.input_mask_torch = F.interpolate(self.input_mask_torch, (self.opt.ref_size, self.opt.ref_size),
                                                  mode="bilinear", align_corners=False)
            self.input_img_torch_channel_last = self.input_img_torch[0].permute(1, 2, 0).contiguous()
        '''

        # TODO: self.input_img_torch_list  input_mask_torch_list
        print('input_img_list:', len(self.input_img_list))
        for img_idx in range(len(self.input_img_list)):
            input_img_torch = torch.from_numpy(self.input_img_list[img_idx]).permute(2, 0, 1).unsqueeze(0).to(self.device)
            input_img_torch = F.interpolate(input_img_torch, (self.opt.ref_size, self.opt.ref_size),
                                            mode="bilinear", align_corners=False)
            self.input_img_torch_list.append(input_img_torch)

            input_mask_torch = torch.from_numpy(self.input_mask_list[img_idx]).permute(2, 0, 1).unsqueeze(0).to(self.device)
            input_mask_torch = F.interpolate(input_mask_torch, (self.opt.ref_size, self.opt.ref_size),
                                             mode="bilinear", align_corners=False)
            self.input_mask_torch_list.append(input_mask_torch)
            self.input_img_torch_channel_last_list.append(input_img_torch[0].permute(1, 2, 0).contiguous())

        # prepare embeddings
        with torch.no_grad():
            # TODO: self.input_img_torch_list  just use first image calc embedding because utils train_step not handle multi
            # prompts negative_prompts support multi?but encode_text in utils not support?
            # support multi image embedding in future? necessary?
            input_img_torch = self.input_img_torch_list[0] if (len(self.input_img_torch_list) > 0) else None
            if self.enable_sd:
                if self.opt.imagedream:  # NOTICE: mvdream only use text input, imagedream both?
                    self.guidance_sd.get_image_text_embeds(input_img_torch, [self.prompt], [self.negative_prompt])
                else:  # only use text input
                    self.guidance_sd.get_text_embeds([self.prompt], [self.negative_prompt])

            if self.enable_zero123:  # only use image input
                self.guidance_zero123.get_img_embeds(input_img_torch)

    def train_step(self):
        starter = torch.cuda.Event(enable_timing=True)
        ender = torch.cuda.Event(enable_timing=True)
        starter.record()

        for step in range(self.train_steps):  # 1??
            print('step:', step)
            self.step += 1
            step_ratio = min(1, self.step / self.opt.iters_refine)

            loss = 0

            ### known view has ground truth
            if not self.opt.imagedream:  # NOTICE: imagedream not use as ground truth?
                print('known view ground truth loss...')
                for idx_image in range(len(self.input_img_torch_list)):
                    # TODO: set cur_cam
                    #cur_cam = self.fixed_cam
                    cur_cam = self.fixed_cam_list[idx_image]

                    ssaa = min(2.0, max(0.125, 2 * np.random.random()))
                    out = self.renderer.render(*cur_cam, self.opt.ref_size, self.opt.ref_size, ssaa=ssaa)

                    # rgb loss
                    image = out["image"]  # [H, W, 3] in [0, 1]
                    valid_mask = ((out["alpha"] > 0) & (out["viewcos"] > 0.5)).detach()
                    mse_loss_rgb = F.mse_loss(image * valid_mask, self.input_img_torch_channel_last_list[idx_image] * valid_mask)
                    print('mse_loss_rgb:', mse_loss_rgb)
                    loss = loss + mse_loss_rgb

            ### novel view (manual batch) no ground truth so use guidance loss
            print('novel view guidance loss...')
            render_resolution = 512
            images = []
            poses = []
            vers, hors, radii = [], [], []
            # avoid too large elevation (> 80 or < -80), and make sure it always cover [min_ver, max_ver]
            min_ver = max(min(self.opt.min_ver, self.opt.min_ver - self.opt.elevation), -80 - self.opt.elevation)
            max_ver = min(max(self.opt.max_ver, self.opt.max_ver - self.opt.elevation), 80 - self.opt.elevation)
            print('batch_size:', self.opt.batch_size)
            for batch_i in range(self.opt.batch_size):
                print('batch_i:', batch_i)
                # render random view
                ver = np.random.randint(min_ver, max_ver)
                hor = np.random.randint(-180, 180)
                radius = 0

                vers.append(ver)
                hors.append(hor)
                radii.append(radius)

                pose = orbit_camera(self.opt.elevation + ver, hor, self.opt.radius + radius)
                # print('pose:', pose.shape)
                # print('pose after orbit_camera', pose)
                poses.append(pose)

                # random render resolution
                ssaa = min(2.0, max(0.125, 2 * np.random.random()))
                print('render ...', batch_i)
                out = self.renderer.render(pose, self.cam.perspective, render_resolution, render_resolution, ssaa=ssaa)
                print('render finish!', batch_i)
                image = out["image"]  # [H, W, 3] in [0, 1]
                image = image.permute(2, 0, 1).contiguous().unsqueeze(0)  # [1, 3, H, W] in [0, 1]

                images.append(image)

                # enable mvdream training
                if self.opt.mvdream or self.opt.imagedream:  # other views: rotate horizontally based random view
                    print('mvdream or imagedream...')
                    for view_i in range(1, 4):  # 3 views each rotate 90deg
                        pose_i = orbit_camera(self.opt.elevation + ver, hor + 90 * view_i, self.opt.radius + radius)
                        print('pose_i:', pose_i.shape)
                        poses.append(pose_i)

                        out_i = self.renderer.render(pose_i, self.cam.perspective, render_resolution, render_resolution,
                                                     ssaa=ssaa)

                        image = out_i["image"].permute(2, 0, 1).contiguous().unsqueeze(0)  # [1, 3, H, W] in [0, 1]
                        images.append(image)

            images = torch.cat(images, dim=0)
            poses = torch.from_numpy(np.stack(poses, axis=0)).to(self.device)

            # import kiui
            # kiui.lo(hor, ver)
            # kiui.vis.plot_image(image)

            # guidance loss
            print('guidance loss start')
            # TODO:
            #  all has embeddings which inited in prepare_train should handle multi images or texts
            #  mvdream:text imagedream image&text sd:text zero123:image
            #  only zero123 has cam_embeddings for novel views
            strength = step_ratio * 0.15 + 0.8
            if self.enable_sd:
                print('guidance_sd.train_step...')
                if self.opt.mvdream or self.opt.imagedream:
                    # loss = loss + self.opt.lambda_sd * self.guidance_sd.train_step(images, poses, step_ratio)
                    refined_images = self.guidance_sd.refine(images, poses, strength=strength).float()
                    refined_images = F.interpolate(refined_images, (render_resolution, render_resolution),
                                                   mode="bilinear", align_corners=False)
                    loss = loss + self.opt.lambda_sd * F.mse_loss(images, refined_images)
                else:
                    # loss = loss + self.opt.lambda_sd * self.guidance_sd.train_step(images, step_ratio)
                    refined_images = self.guidance_sd.refine(images, strength=strength).float()
                    refined_images = F.interpolate(refined_images, (render_resolution, render_resolution),
                                                   mode="bilinear", align_corners=False)
                    sd_loss = F.mse_loss(images, refined_images)
                    print('sd_loss:', sd_loss)
                    loss = loss + self.opt.lambda_sd * sd_loss
                print('guidance_sd.train_step finish!')
            if self.enable_zero123:
                print('guidance_zero123.train_step ...')
                # loss = loss + self.opt.lambda_zero123 * self.guidance_zero123.train_step(images, vers, hors, radii, step_ratio)
                refined_images = self.guidance_zero123.refine(images, vers, hors, radii, strength=strength,
                                                              default_elevation=self.opt.elevation).float()
                refined_images = F.interpolate(refined_images, (render_resolution, render_resolution), mode="bilinear",
                                               align_corners=False)
                sd_loss = F.mse_loss(images, refined_images)
                print('sd_loss:', sd_loss)
                loss = loss + self.opt.lambda_zero123 * sd_loss
                # loss = loss + self.opt.lambda_zero123 * self.lpips_loss(images, refined_images)
                print('guidance_zero123.train_step finish!')
            # optimize step
            print('optimize step ...')
            loss.backward()
            self.optimizer.step()
            self.optimizer.zero_grad()
            print('optimize step finish!')

        ender.record()
        torch.cuda.synchronize()
        t = starter.elapsed_time(ender)

        self.need_update = True

        if self.gui:
            dpg.set_value("_log_train_time", f"{t:.4f}ms")
            dpg.set_value(
                "_log_train_log",
                f"step = {self.step: 5d} (+{self.train_steps: 2d}) loss = {loss.item():.4f}",
            )

        # dynamic train steps (no need for now)
        # max allowed train time per-frame is 500 ms
        # full_t = t / self.train_steps * 16
        # train_steps = min(16, max(4, int(16 * 500 / full_t)))
        # if train_steps > self.train_steps * 1.2 or train_steps < self.train_steps * 0.8:
        #     self.train_steps = train_steps

    @torch.no_grad()
    def test_step(self):
        # ignore if no need to update
        if not self.need_update:
            return

        starter = torch.cuda.Event(enable_timing=True)
        ender = torch.cuda.Event(enable_timing=True)
        starter.record()

        # should update image
        if self.need_update:
            # render image

            out = self.renderer.render(self.cam.pose, self.cam.perspective, self.H, self.W)

            buffer_image = out[self.mode]  # [H, W, 3]

            if self.mode in ['depth', 'alpha']:
                buffer_image = buffer_image.repeat(1, 1, 3)
                if self.mode == 'depth':
                    buffer_image = (buffer_image - buffer_image.min()) / (
                            buffer_image.max() - buffer_image.min() + 1e-20)

            self.buffer_image = buffer_image.contiguous().clamp(0, 1).detach().cpu().numpy()

            # display input_image
            if self.overlay_input_img and self.input_img is not None:
                self.buffer_image = (
                        self.buffer_image * (1 - self.overlay_input_img_ratio)
                        + self.input_img * self.overlay_input_img_ratio
                )

            self.need_update = False

        ender.record()
        torch.cuda.synchronize()
        t = starter.elapsed_time(ender)

        if self.gui:
            dpg.set_value("_log_infer_time", f"{t:.4f}ms ({int(1000 / t)} FPS)")
            dpg.set_value(
                "_texture", self.buffer_image
            )  # buffer must be contiguous, else seg fault!

    def load_input(self, file):
        print('file:', file)
        self.input_img, self.input_mask = self.load_image(file)
        # print(self.input_img, self.input_mask)
        # exit(0)
        self.prompt = self.load_prompt(file)
        print('prompt：', self.prompt)

    def load_input_files(self, files):
        for file in files:
            print('file:', file)
            input_img, input_mask = self.load_image(file)
            self.input_img_list.append(input_img)
            self.input_mask_list.append(input_mask)
            prompt = self.load_prompt(file)
            print('prompt：', prompt)
            self.prompt_list.append(prompt)

    def load_image(self, file):
        # load image
        print(f'[INFO] load image from {file} ...')
        img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
        print('img shape:', img.shape)
        if img.shape[-1] == 3:
            if self.bg_remover is None:
                self.bg_remover = rembg.new_session()
            img = rembg.remove(img, session=self.bg_remover)

        img = cv2.resize(img, (self.W, self.H), interpolation=cv2.INTER_AREA)
        img = img.astype(np.float32) / 255.0

        input_mask = img[..., 3:]
        # white bg
        input_img = img[..., :3] * input_mask + (1 - input_mask)
        # bgr to rgb
        input_img = input_img[..., ::-1].copy()
        return input_img, input_mask

    def load_prompt(self, file):
        # load prompt
        prompt = ""
        file_prompt = file.replace("_rgba.png", "_caption.txt")
        if os.path.exists(file_prompt):
            print(f'[INFO] load prompt from {file_prompt}...')
            with open(file_prompt, "r") as f:
                prompt = f.read().strip()
        return prompt

    def save_model(self):
        os.makedirs(self.opt.outdir, exist_ok=True)

        path = os.path.join(self.opt.outdir, self.opt.save_path + '.' + self.opt.mesh_format)
        self.renderer.export_mesh(path)

        print(f"[INFO] save model to {path} .")

    def register_dpg(self):
        ### register texture

        with dpg.texture_registry(show=False):
            dpg.add_raw_texture(
                self.W,
                self.H,
                self.buffer_image,
                format=dpg.mvFormat_Float_rgb,
                tag="_texture",
            )

        ### register window

        # the rendered image, as the primary window
        with dpg.window(
                tag="_primary_window",
                width=self.W,
                height=self.H,
                pos=[0, 0],
                no_move=True,
                no_title_bar=True,
                no_scrollbar=True,
        ):
            # add the texture
            dpg.add_image("_texture")

        # dpg.set_primary_window("_primary_window", True)

        # control window
        with dpg.window(
                label="Control",
                tag="_control_window",
                width=600,
                height=self.H,
                pos=[self.W, 0],
                no_move=True,
                no_title_bar=True,
        ):
            # button theme
            with dpg.theme() as theme_button:
                with dpg.theme_component(dpg.mvButton):
                    dpg.add_theme_color(dpg.mvThemeCol_Button, (23, 3, 18))
                    dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (51, 3, 47))
                    dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (83, 18, 83))
                    dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 5)
                    dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 3, 3)

            # timer stuff
            with dpg.group(horizontal=True):
                dpg.add_text("Infer time: ")
                dpg.add_text("no data", tag="_log_infer_time")

            def callback_setattr(sender, app_data, user_data):
                setattr(self, user_data, app_data)

            # init stuff
            with dpg.collapsing_header(label="Initialize", default_open=True):

                # seed stuff
                def callback_set_seed(sender, app_data):
                    self.seed = app_data
                    self.seed_everything()

                dpg.add_input_text(
                    label="seed",
                    default_value=self.seed,
                    on_enter=True,
                    callback=callback_set_seed,
                )

                # input stuff
                def callback_select_input(sender, app_data):
                    # only one item
                    for k, v in app_data["selections"].items():
                        dpg.set_value("_log_input", k)
                        self.load_input(v)

                    self.need_update = True

                with dpg.file_dialog(
                        directory_selector=False,
                        show=False,
                        callback=callback_select_input,
                        file_count=1,
                        tag="file_dialog_tag",
                        width=700,
                        height=400,
                ):
                    dpg.add_file_extension("Images{.jpg,.jpeg,.png}")

                with dpg.group(horizontal=True):
                    dpg.add_button(
                        label="input",
                        callback=lambda: dpg.show_item("file_dialog_tag"),
                    )
                    dpg.add_text("", tag="_log_input")

                # overlay stuff
                with dpg.group(horizontal=True):
                    def callback_toggle_overlay_input_img(sender, app_data):
                        self.overlay_input_img = not self.overlay_input_img
                        self.need_update = True

                    dpg.add_checkbox(
                        label="overlay image",
                        default_value=self.overlay_input_img,
                        callback=callback_toggle_overlay_input_img,
                    )

                    def callback_set_overlay_input_img_ratio(sender, app_data):
                        self.overlay_input_img_ratio = app_data
                        self.need_update = True

                    dpg.add_slider_float(
                        label="ratio",
                        min_value=0,
                        max_value=1,
                        format="%.1f",
                        default_value=self.overlay_input_img_ratio,
                        callback=callback_set_overlay_input_img_ratio,
                    )

                # prompt stuff

                dpg.add_input_text(
                    label="prompt",
                    default_value=self.prompt,
                    callback=callback_setattr,
                    user_data="prompt",
                )

                dpg.add_input_text(
                    label="negative",
                    default_value=self.negative_prompt,
                    callback=callback_setattr,
                    user_data="negative_prompt",
                )

                # save current model
                with dpg.group(horizontal=True):
                    dpg.add_text("Save: ")

                    dpg.add_button(
                        label="model",
                        tag="_button_save_model",
                        callback=self.save_model,
                    )
                    dpg.bind_item_theme("_button_save_model", theme_button)

                    dpg.add_input_text(
                        label="",
                        default_value=self.opt.save_path,
                        callback=callback_setattr,
                        user_data="save_path",
                    )

            # training stuff
            with dpg.collapsing_header(label="Train", default_open=True):
                # lr and train button
                with dpg.group(horizontal=True):
                    dpg.add_text("Train: ")

                    def callback_train(sender, app_data):
                        if self.training:
                            self.training = False
                            dpg.configure_item("_button_train", label="start")
                        else:
                            self.prepare_train()
                            self.training = True
                            dpg.configure_item("_button_train", label="stop")

                    # dpg.add_button(
                    #     label="init", tag="_button_init", callback=self.prepare_train
                    # )
                    # dpg.bind_item_theme("_button_init", theme_button)

                    dpg.add_button(
                        label="start", tag="_button_train", callback=callback_train
                    )
                    dpg.bind_item_theme("_button_train", theme_button)

                with dpg.group(horizontal=True):
                    dpg.add_text("", tag="_log_train_time")
                    dpg.add_text("", tag="_log_train_log")

            # rendering options
            with dpg.collapsing_header(label="Rendering", default_open=True):
                # mode combo
                def callback_change_mode(sender, app_data):
                    self.mode = app_data
                    self.need_update = True

                dpg.add_combo(
                    ("image", "depth", "alpha", "normal"),
                    label="mode",
                    default_value=self.mode,
                    callback=callback_change_mode,
                )

                # fov slider
                def callback_set_fovy(sender, app_data):
                    self.cam.fovy = np.deg2rad(app_data)
                    self.need_update = True

                dpg.add_slider_int(
                    label="FoV (vertical)",
                    min_value=1,
                    max_value=120,
                    format="%d deg",
                    default_value=np.rad2deg(self.cam.fovy),
                    callback=callback_set_fovy,
                )

        ### register camera handler

        def callback_camera_drag_rotate_or_draw_mask(sender, app_data):
            if not dpg.is_item_focused("_primary_window"):
                return

            dx = app_data[1]
            dy = app_data[2]

            self.cam.orbit(dx, dy)
            self.need_update = True

        def callback_camera_wheel_scale(sender, app_data):
            if not dpg.is_item_focused("_primary_window"):
                return

            delta = app_data

            self.cam.scale(delta)
            self.need_update = True

        def callback_camera_drag_pan(sender, app_data):
            if not dpg.is_item_focused("_primary_window"):
                return

            dx = app_data[1]
            dy = app_data[2]

            self.cam.pan(dx, dy)
            self.need_update = True

        def callback_set_mouse_loc(sender, app_data):
            if not dpg.is_item_focused("_primary_window"):
                return

            # just the pixel coordinate in image
            self.mouse_loc = np.array(app_data)

        with dpg.handler_registry():
            # for camera moving
            dpg.add_mouse_drag_handler(
                button=dpg.mvMouseButton_Left,
                callback=callback_camera_drag_rotate_or_draw_mask,
            )
            dpg.add_mouse_wheel_handler(callback=callback_camera_wheel_scale)
            dpg.add_mouse_drag_handler(
                button=dpg.mvMouseButton_Middle, callback=callback_camera_drag_pan
            )

        dpg.create_viewport(
            title="Gaussian3D",
            width=self.W + 600,
            height=self.H + (45 if os.name == "nt" else 0),
            resizable=False,
        )

        ### global theme
        with dpg.theme() as theme_no_padding:
            with dpg.theme_component(dpg.mvAll):
                # set all padding to 0 to avoid scroll bar
                dpg.add_theme_style(
                    dpg.mvStyleVar_WindowPadding, 0, 0, category=dpg.mvThemeCat_Core
                )
                dpg.add_theme_style(
                    dpg.mvStyleVar_FramePadding, 0, 0, category=dpg.mvThemeCat_Core
                )
                dpg.add_theme_style(
                    dpg.mvStyleVar_CellPadding, 0, 0, category=dpg.mvThemeCat_Core
                )

        dpg.bind_item_theme("_primary_window", theme_no_padding)

        dpg.setup_dearpygui()

        ### register a larger font
        # get it from: https://github.com/lxgw/LxgwWenKai/releases/download/v1.300/LXGWWenKai-Regular.ttf
        if os.path.exists("LXGWWenKai-Regular.ttf"):
            with dpg.font_registry():
                with dpg.font("LXGWWenKai-Regular.ttf", 18) as default_font:
                    dpg.bind_font(default_font)

        # dpg.show_metrics()

        dpg.show_viewport()

    def render(self):
        assert self.gui
        while dpg.is_dearpygui_running():
            # update texture every frame
            if self.training:
                print('train_step ...')
                self.train_step()
                print('train_step finish!')
            print('test_step ...')
            self.test_step()
            print('test_step finish!')
            dpg.render_dearpygui_frame()
            print('render_dearpygui_frame finish!')

    # no gui mode
    def train(self, iters=500):
        if iters > 0:
            print('prepare_train ...')
            self.prepare_train()
            print('prepare_train finish!')
            for i in tqdm.trange(iters):
                self.train_step()
            print('train_steps finish!iters:', iters)
        # save
        print('save_model model ...')
        self.save_model()
        print('save_model model finish!')


if __name__ == "__main__":
    import argparse
    from omegaconf import OmegaConf

    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="path to the yaml config file")
    args, extras = parser.parse_known_args()
    # print('args:', args, 'extras:', extras)
    # exit(0)
    # override default config from cli
    opt = OmegaConf.merge(OmegaConf.load(args.config), OmegaConf.from_cli(extras))

    # auto find mesh from stage 1
    if opt.mesh is None:
        default_path = os.path.join(opt.outdir, opt.save_path + '_mesh.' + opt.mesh_format)
        if os.path.exists(default_path):
            opt.mesh = default_path
        else:
            raise ValueError(f"Cannot find mesh from {default_path}, must specify --mesh explicitly!")

    gui = GUI(opt)

    if opt.gui:
        print('gui render')
        gui.render()
    else:
        print('gui train')
        gui.train(opt.iters_refine)