pipeline/kiss3d_wrapper.py

# The kiss3d pipeline wrapper for inference

import os
import numpy as np
import torch
import yaml
import uuid
from typing import Union, Any, Dict
from einops import rearrange
from PIL import Image

from pipeline.utils import logger, TMP_DIR, OUT_DIR
from pipeline.utils import lrm_reconstruct, isomer_reconstruct

import torch
import torchvision

# for reconstruction model
from omegaconf import OmegaConf
from models.lrm.utils.train_util import instantiate_from_config
from models.lrm.utils.render_utils import rotate_x, rotate_y
from utils.tool import get_background

# for florence2
from transformers import AutoProcessor, AutoModelForCausalLM

from diffusers import FluxPipeline, FluxControlNetImg2ImgPipeline, FluxImg2ImgPipeline, DiffusionPipeline, EulerAncestralDiscreteScheduler
from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel


def init_wrapper_from_config(config_path):
    with open(config_path, 'r') as config_file:
        config_ = yaml.load(config_file, yaml.FullLoader)
    
    # init flux_pipeline
    logger.info('==> Loading Flux model ...')
    flux_device = config_['flux'].get('device', 'cpu')
    flux_base_model_pth = config_['flux'].get('base_model', None)
    flux_controlnet_pth = config_['flux'].get('controlnet', None)
    flux_lora_pth = config_['flux'].get('lora', None)

    # load flux model and controlnet
    if flux_controlnet_pth is not None:
        flux_controlnet = FluxControlNetModel.from_pretrained(flux_controlnet_pth)
        flux_pipe = FluxControlNetImg2ImgPipeline.from_pretrained(flux_base_model_pth, controlnet=[flux_controlnet], \
                                torch_dtype=torch.bfloat16)
    else:
        flux_pipe = FluxImg2ImgPipeline(flux_base_model_pth, torch_dtype=torch.bfloat16)
    
    # load lora weights
    flux_pipe.load_lora_weights(flux_lora_pth)
    flux_pipe.to(device=flux_device, dtype=torch.bfloat16)

    # TODO: load redux model
    # FluxPriorReduxPipeline.from_pretrained()

    # TODO: load pulid model

    # init multiview model
    logger.info('==> Loading multiview diffusion model ...')
    multiview_device = config_['multiview'].get('device', 'cpu')
    multiview_pipeline = DiffusionPipeline.from_pretrained(
        config_['multiview']['base_model'], 
        custom_pipeline=config_['multiview']['custom_pipeline'],
        torch_dtype=torch.float16,
    )
    multiview_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
        multiview_pipeline.scheduler.config, timestep_spacing='trailing'
    )
    
    unet_ckpt_path = config_['multiview'].get('unet', None)
    if unet_ckpt_path is not None:
        state_dict = torch.load(unet_ckpt_path, map_location='cpu')['state_dict']
        state_dict = {k[10:]: v for k, v in state_dict.items() if k.startswith('unet.unet.')}
        multiview_pipeline.unet.load_state_dict(state_dict, strict=True)

    multiview_pipeline.to(multiview_device)

    # load caption model
    logger.info('==> Loading caption model ...')
    caption_device = config_['caption'].get('device', 'cpu')
    caption_model = AutoModelForCausalLM.from_pretrained(config_['caption']['base_model'], \
                    torch_dtype=torch.bfloat16, trust_remote_code=True).to(caption_device)
    caption_processor = AutoProcessor.from_pretrained(config_['caption']['base_model'], trust_remote_code=True)

    # load reconstruction model
    logger.info('==> Loading reconstruction model ...')
    recon_device = config_['reconstruction'].get('device', 'cpu')
    recon_model_config = OmegaConf.load(config_['reconstruction']['model_config'])
    recon_model = instantiate_from_config(recon_model_config.model_config)
    # load recon model checkpoint
    state_dict = torch.load(config_['reconstruction']['base_model'], map_location='cpu')['state_dict']
    state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
    recon_model.load_state_dict(state_dict, strict=True)
    recon_model.to(recon_device)
    recon_model.init_flexicubes_geometry(recon_device, fovy=50.0)
    recon_model.eval()

    return kiss3d_wrapper(
        config = config_,
        flux_pipeline = flux_pipe,
        multiview_pipeline = multiview_pipeline,
        caption_processor = caption_processor,
        caption_model = caption_model,
        reconstruction_model_config = recon_model_config,
        reconstruction_model = recon_model,
    )

class kiss3d_wrapper(object):
    def __init__(self,
        config: Dict,
        flux_pipeline: Union[FluxPipeline, FluxControlNetImg2ImgPipeline],
        multiview_pipeline: DiffusionPipeline,
        caption_processor: AutoProcessor,
        caption_model: AutoModelForCausalLM,
        reconstruction_model_config: Any,
        reconstruction_model: Any,
    ):
        self.config = config
        self.flux_pipeline = flux_pipeline
        self.multiview_pipeline = multiview_pipeline
        self.caption_model = caption_model
        self.caption_processor = caption_processor
        self.recon_model_config = reconstruction_model_config
        self.recon_model = reconstruction_model        

        self.renew_uuid()

    def renew_uuid(self):
        self.uuid = uuid.uuid4()

    def context(self):
        if self.config['use_zero_gpu']:
            import spaces
            return spaces.GPU()
        else:
            return torch.no_grad()

    def get_image_caption(self, image):
        """
        image: PIL image or path of PIL image
        """
        torch_dtype = torch.bfloat16
        caption_device = self.config['caption'].get('device', 'cpu')

        if isinstance(image, str):  # If image is a file path
            image = Image.open(image).convert("RGB")
        elif isinstance(image, Image):
            image = image.convert("RGB")
        else:
            raise NotImplementedError('unexpected image type')
        
        prompt = "<MORE_DETAILED_CAPTION>"
        inputs = self.caption_processor(text=prompt, images=image, return_tensors="pt").to(caption_device, torch_dtype)

        generated_ids = self.caption_model.generate(
                input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3
            )

        generated_text = self.caption_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
        parsed_answer = self.caption_processor.post_process_generation(
            generated_text, task=prompt, image_size=(image.width, image.height)
        )
        caption_text = parsed_answer["<MORE_DETAILED_CAPTION>"].replace("The image is ", "")
        return caption_text

    def generate_multiview(self, image):
        with self.context():
            mv_image = self.multiview_pipeline(image, 
                                               num_inference_steps=self.config['multiview']['num_inference_steps'], 
                                               width=512*2, height=512*2).images[0]
        return mv_image

    def reconstruct_from_multiview(self, mv_image):
        """
        mv_image: PIL.Image
        """
        recon_device = self.config['reconstruction'].get('device', 'cpu')

        rgb_multi_view = np.asarray(mv_image, dtype=np.float32) / 255.0
        rgb_multi_view = torch.from_numpy(rgb_multi_view).squeeze(0).permute(2, 0, 1).contiguous().float()     # (3, 1024, 2048)
        rgb_multi_view = rearrange(rgb_multi_view, 'c (n h) (m w) -> (n m) c h w', n=2, m=2).unsqueeze(0).to(recon_device)

        with self.context():
            vertices, faces, lrm_multi_view_normals, lrm_multi_view_rgb, lrm_multi_view_albedo = \
            lrm_reconstruct(self.recon_model, self.recon_model_config.infer_config,
                            rgb_multi_view, name=self.uuid)

        return vertices, faces, lrm_multi_view_normals, lrm_multi_view_rgb, lrm_multi_view_albedo
    
    def generate_reference_3D_bundle_image_zero123(self, image, save_intermediate_results=True):
        """
        input: image, PIL.Image
        return: ref_3D_bundle_image, Tensor of shape (1, 3, 1024, 2048)
        """
        mv_image = self.generate_multiview(image)

        if save_intermediate_results:
            mv_image.save(os.path.join(TMP_DIR, f'{self.uuid}_mv_image.png'))

        vertices, faces, lrm_multi_view_normals, lrm_multi_view_rgb, lrm_multi_view_albedo = self.reconstruct_from_multiview(mv_image)

        ref_3D_bundle_image = torchvision.utils.make_grid(torch.cat([lrm_multi_view_rgb.cpu(), (lrm_multi_view_normals.cpu() + 1) / 2], dim=0), nrow=4, padding=0).unsqueeze(0) # range [0, 1]
        
        if save_intermediate_results:
            save_path = os.path.join(TMP_DIR, f'{self.uuid}_ref_3d_bundle_image.png')
            torchvision.utils.save_image(ref_3D_bundle_image, save_path)

            logger.info(f"Save reference 3D bundle image to {save_path}")

            return ref_3D_bundle_image, save_path

        return ref_3D_bundle_image

    def generate_3d_bundle_image_controlnet(self, 
                                 prompt, 
                                 image=None,
                                 strength=1.0, 
                                 control_image=[],
                                 control_mode=[],
                                 control_guidance_start=None,
                                 control_guidance_end=None,
                                 controlnet_conditioning_scale=None,
                                 lora_scale=1.0,
                                 save_intermediate_results=True,
                                 **kwargs):
        control_mode_dict = {
            'canny': 0,
            'tile': 1,
            'depth': 2,
            'blur': 3,
            'pose': 4,
            'gray': 5,
            'lq': 6,
        } # for https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union only

        flux_device = self.config['flux'].get('device', 'cpu')
        seed = self.config['flux'].get('seed', 0)

        generator = torch.Generator(device=flux_device).manual_seed(seed)

        hparam_dict = {
            'prompt': ' '.join(['A grid of 2x4 multi-view image, elevation 5. White background.', prompt]),
            'image': image or torch.zeros((1, 3, 1024, 2048), dtype=torch.float32, device=flux_device),
            'strength': strength,
            'num_inference_steps': 30,
            'guidance_scale': 3.5,
            'num_images_per_prompt': 1,
            'width': 2048,
            'height': 1024,
            'output_type': 'np',
            'generator': generator,
            'joint_attention_kwargs': {"scale": lora_scale}
        }
        hparam_dict.update(kwargs)
        
         # append controlnet hparams
        if len(control_image) > 0:
            assert isinstance(self.flux_pipeline, FluxControlNetImg2ImgPipeline)
            assert len(control_mode) == len(control_image) # the count of image should be the same as control mode
            
            flux_ctrl_net = self.flux_pipeline.controlnet.nets[0]
            self.flux_pipeline.controlnet = FluxMultiControlNetModel([flux_ctrl_net for i in range(len(control_image))])

            ctrl_hparams = {
                'control_mode': [control_mode_dict[mode_] for mode_ in control_mode],
                'control_image': control_image,
                'control_guidance_start': control_guidance_start or [0.0 for i in range(len(control_image))],
                'control_guidance_end': control_guidance_end or [1.0 for i in range(len(control_image))],
                'controlnet_conditioning_scale': controlnet_conditioning_scale or [1.0 for i in range(len(control_image))],
            }

            hparam_dict.update(ctrl_hparams)

        with self.context():
            gen_3d_bundle_image = self.flux_pipeline(**hparam_dict).images
        
        gen_3d_bundle_image_ = torch.from_numpy(gen_3d_bundle_image).squeeze(0).permute(2, 0, 1).contiguous().float()     # (3, 1024, 2048)

        if save_intermediate_results:
            save_path = os.path.join(TMP_DIR, f'{self.uuid}_gen_3d_bundle_image.png')
            torchvision.utils.save_image(gen_3d_bundle_image_, save_path)
            logger.info(f"Save generated 3D bundle image to {save_path}")
            return gen_3d_bundle_image_, save_path

        return gen_3d_bundle_image_


    def generate_3d_bundle_image_text(self, 
                                      prompt,
                                      image=None, 
                                      strength=1.0,
                                      lora_scale=1.0,
                                      num_inference_steps=30,
                                      save_intermediate_results=True,
                                      **kwargs):
        
        """
        return: gen_3d_bundle_image, torch.Tensor of shape (3, 1024, 2048), range [0., 1.]
        """
        
        if isinstance(self.flux_pipeline, FluxControlNetImg2ImgPipeline):
            flux_pipeline = FluxImg2ImgPipeline(
                scheduler = self.flux_pipeline.scheduler,
                vae = self.flux_pipeline.vae,
                text_encoder = self.flux_pipeline.text_encoder,
                tokenizer = self.flux_pipeline.tokenizer,
                text_encoder_2 = self.flux_pipeline.text_encoder_2,
                tokenizer_2 = self.flux_pipeline.tokenizer_2,
                transformer = self.flux_pipeline.transformer
            )
        else:
            flux_pipeline = self.flux_pipeline

        flux_device = self.config['flux'].get('device', 'cpu')
        seed = self.config['flux'].get('seed', 0)

        generator = torch.Generator(device=flux_device).manual_seed(seed)

        hparam_dict = {
            'prompt': ' '.join(['A grid of 2x4 multi-view image, elevation 5. White background.', prompt]),
            'image': image or torch.zeros((1, 3, 1024, 2048), dtype=torch.float32, device=flux_device),
            'strength': strength,
            'num_inference_steps': num_inference_steps,
            'guidance_scale': 3.5,
            'num_images_per_prompt': 1,
            'width': 2048,
            'height': 1024,
            'output_type': 'np',
            'generator': generator,
            'joint_attention_kwargs': {"scale": lora_scale}
        }
        hparam_dict.update(kwargs)

        with self.context():
            gen_3d_bundle_image = flux_pipeline(**hparam_dict).images

        gen_3d_bundle_image_ = torch.from_numpy(gen_3d_bundle_image).squeeze(0).permute(2, 0, 1).contiguous().float()     # (3, 1024, 2048)

        if save_intermediate_results:
            save_path = os.path.join(TMP_DIR, f'{self.uuid}_gen_3d_bundle_image.png')
            torchvision.utils.save_image(gen_3d_bundle_image_, save_path)
            logger.info(f"Save generated 3D bundle image to {save_path}")
            return gen_3d_bundle_image_, save_path

        return gen_3d_bundle_image_
    
    def reconstruct_3d_bundle_image(self, image, save_intermediate_results=True):
        """
        image: torch.Tensor, range [0., 1.], (3, 1024, 2048)
        """
        recon_device = self.config['reconstruction'].get('device', 'cpu')

        # split rgb and normal
        images = rearrange(image, 'c (n h) (m w) -> (n m) c h w', n=2, m=4) # (3, 1024, 2048) -> (8, 3, 512, 512)
        rgb_multi_view, normal_multi_view = images.chunk(2, dim=0)
        multi_view_mask = get_background(normal_multi_view).to(recon_device)
        rgb_multi_view = rgb_multi_view.to(recon_device) * multi_view_mask + (1 - multi_view_mask)
        
        with self.context():
            vertices, faces, lrm_multi_view_normals, lrm_multi_view_rgb, lrm_multi_view_albedo = \
            lrm_reconstruct(self.recon_model, self.recon_model_config.infer_config,
                            rgb_multi_view.unsqueeze(0).to(recon_device), name=self.uuid, 
                            input_camera_type='kiss3d', render_3d_bundle_image=save_intermediate_results,
                            render_azimuths=[0, 90, 180, 270])

        if save_intermediate_results:
            recon_3D_bundle_image = torchvision.utils.make_grid(torch.cat([lrm_multi_view_rgb.cpu(), (lrm_multi_view_normals.cpu() + 1) / 2], dim=0), nrow=4, padding=0).unsqueeze(0) # range [0, 1]        
            torchvision.utils.save_image(recon_3D_bundle_image, os.path.join(TMP_DIR, f'{k3d_wrapper.uuid})_lrm_recon_3d_bundle_image.png'))

        recon_mesh_path = os.path.join(TMP_DIR, f"{self.uuid}_isomer_recon_mesh.obj")
        
        return isomer_reconstruct(rgb_multi_view=rgb_multi_view,
                                  normal_multi_view=normal_multi_view,
                                  multi_view_mask=multi_view_mask,
                                  vertices=vertices,
                                  faces=faces,
                                  save_path=recon_mesh_path)


def run_text_to_3d(k3d_wrapper,
                   prompt,
                   init_image_path=None):
    # ======================================= Example of text to 3D generation ======================================

    # Renew The uuid
    k3d_wrapper.renew_uuid()

    # FOR Text to 3D (also for image to image) with init image
    init_image = None
    if init_image_path is not None:
        init_image = Image.open(init_image_path)

    gen_3d_bundle_image, gen_save_path = k3d_wrapper.generate_3d_bundle_image_text(prompt, 
                                                                                     image=init_image, 
                                                                                     strength=1.0, 
                                                                                     save_intermediate_results=True)

    # recon from 3D Bundle image
    recon_mesh_path = k3d_wrapper.reconstruct_3d_bundle_image(gen_3d_bundle_image, save_intermediate_results=False)

    return gen_save_path, recon_mesh_path

def run_image_to_3d(k3d_wrapper, init_image_path):
    # ======================================= Example of image to 3D generation ======================================

    # Renew The uuid
    k3d_wrapper.renew_uuid()

    # FOR IMAGE TO 3D: generate reference 3D bundle image from a single input image
    input_image = Image.open(init_image_path)
    reference_3d_bundle_image, reference_save_path = k3d_wrapper.generate_reference_3D_bundle_image_zero123(input_image)
    caption = k3d_wrapper.get_image_caption(input_image)


    import pdb
    pdb.set_trace()


if __name__ == "__main__":
    k3d_wrapper = init_wrapper_from_config('/hpc2hdd/home/jlin695/code/Kiss3DGen/pipeline/pipeline_config/default.yaml')

    # Example of loading existing 3D bundle Image
    # demo_image = Image.open('/hpc2hdd/home/jlin695/code/github/Kiss3DGen/outputs/tmp/ea25bc9b-d775-46bb-9827-660a9a6540c8_gen_3d_bundle_image.png')
    # gen_3d_bundle_image = torchvision.transforms.functional.to_tensor(demo_image)

    run_image_to_3d(k3d_wrapper, '/hpc2hdd/home/jlin695/code/Kiss3DGen/examples/蓝色小怪物.webp')
    # run_text_to_3d(k3d_wrapper, prompt='A doll of a girl in Harry Potter')