単一画像と OpenVINO から TripoSR フィードフォワード 3D を再構築#

%pip install -q "gradio>=4.19" "torch==2.2.2" rembg trimesh einops "omegaconf>=2.3.0" "transformers>=4.35.0" "openvino>=2024.0.0" --extra-index-url https://download.pytorch.org/whl/cpu 
%pip install -q "git+https://github.com/tatsy/torchmcubes.git"

import sys 
from pathlib import Path 

if not Path("TripoSR").exists():
     !git clone https://huggingface.co/spaces/stabilityai/TripoSR 

sys.path.append(model_repo)

import os 

from tsr.system import TSR 

model = TSR.from_pretrained( 
    "stabilityai/TripoSR", 
    config_name="config.yaml", 
    weight_name="model.ckpt", 
) 
model.renderer.set_chunk_size(131072) 
model.to("cpu")

import ipywidgets as widgets 

core = ov.Core() 
device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

compiled_vit_patch_embeddings = core.compile_model(VIT_PATCH_EMBEDDINGS_OV_PATH, device.value) 
compiled_vit_model_encoder = core.compile_model(VIT_ENCODER_OV_PATH, device.value) 
compiled_vit_model_pooler = core.compile_model(VIT_POOLER_OV_PATH, device.value) 

compiled_tokenizer = core.compile_model(TOKENIZER_OV_PATH, device.value) 
compiled_backbone = core.compile_model(BACKBONE_OV_PATH, device.value) 
compiled_post_processor = core.compile_model(POST_PROCESSOR_OV_PATH, device.value)

from collections import namedtuple 

class VitPatchEmdeddingsWrapper(torch.nn.Module): 
    def __init__(self, vit_patch_embeddings, model): 
        super().__init__() 
        self.vit_patch_embeddings = vit_patch_embeddings 
        self.projection = model.projection 

    def forward(self, pixel_values, interpolate_pos_encoding=False): 
        inputs = { 
            "pixel_values": pixel_values, 
        } 
        outs = self.vit_patch_embeddings(inputs)[0] 

        return torch.from_numpy(outs) 

class VitModelEncoderWrapper(torch.nn.Module): 
    def __init__(self, vit_model_encoder): 
        super().__init__()
         self.vit_model_encoder = vit_model_encoder 

    def forward( 
        self, 
        hidden_states, 
        head_mask, 
        output_attentions=False, 
        output_hidden_states=False, 
        return_dict=False, 
    ): 
        inputs = { 
            "hidden_states": hidden_states.detach().numpy(), 
        } 

        outs = self.vit_model_encoder(inputs) 
        outputs = namedtuple("BaseModelOutput", ("last_hidden_state", "hidden_states", "attentions")) 

        return outputs(torch.from_numpy(outs[0]), None, None) 

class VitModelPoolerWrapper(torch.nn.Module): 
    def __init__(self, vit_model_pooler): 
        super().__init__() 
        self.vit_model_pooler = vit_model_pooler 

    def forward(self, hidden_states): 
        outs = self.vit_model_pooler(hidden_states.detach().numpy())[0] 

        return torch.from_numpy(outs) 

class TokenizerWrapper(torch.nn.Module): 
    def __init__(self, tokenizer, model): 
        super().__init__() 
        self.tokenizer = tokenizer self.detokenize = model.detokenize 

    def forward(self, batch_size): 
        outs = self.tokenizer(batch_size)[0] 

        return torch.from_numpy(outs) 

class BackboneWrapper(torch.nn.Module): 
    def __init__(self, backbone): 
        super().__init__() 
        self.backbone = backbone 

    def forward(self, hidden_states, encoder_hidden_states): 
        inputs = { 
            "hidden_states": hidden_states, 
            "encoder_hidden_states": encoder_hidden_states.detach().numpy(), 
        } 

        outs = self.backbone(inputs)[0] 

        return torch.from_numpy(outs) 

class PostProcessorWrapper(torch.nn.Module): 
    def __init__(self, post_processor): 
        super().__init__() 
        self.post_processor = post_processor 

    def forward(self, triplanes): 
        outs = self.post_processor(triplanes)[0] 

        return torch.from_numpy(outs)

model.image_tokenizer.model.embeddings.patch_embeddings = VitPatchEmdeddingsWrapper( 
    compiled_vit_patch_embeddings, 
    model.image_tokenizer.model.embeddings.patch_embeddings, 
) 
model.image_tokenizer.model.encoder = VitModelEncoderWrapper(compiled_vit_model_encoder) 
model.image_tokenizer.model.pooler = VitModelPoolerWrapper(compiled_vit_model_pooler) 

model.tokenizer = TokenizerWrapper(compiled_tokenizer, model.tokenizer) 
model.backbone = BackboneWrapper(compiled_backbone) 
model.post_processor = PostProcessorWrapper(compiled_post_processor)

import tempfile 

import gradio as gr 
import numpy as np 
import rembg 
from PIL import Image 

from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation 

rembg_session = rembg.new_session() 

def check_input_image(input_image): 
    if input_image is None: 
        raise gr.Error("No image uploaded!") 

def preprocess(input_image, do_remove_background, foreground_ratio): 
    def fill_background(image): 
        image = np.array(image).astype(np.float32) / 255.0 
        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5 
        image = Image.fromarray((image * 255.0).astype(np.uint8)) 
        return image 

    if do_remove_background: 
        image = input_image.convert("RGB") 
        image = remove_background(image, rembg_session) 
        image = resize_foreground(image, foreground_ratio) 
        image = fill_background(image) 
    else: 
        image = input_image 
        if image.mode == "RGBA": 
            image = fill_background(image) 
    return image 

def generate(image): 
    scene_codes = model(image, "cpu") # このデバイスは画像処理装置用に提供されます 
    mesh = model.extract_mesh(scene_codes)[0] 
    mesh = to_gradio_3d_orientation(mesh) 
    mesh_path = tempfile.NamedTemporaryFile(suffix=".obj", delete=False) 
    mesh.export(mesh_path.name) 
    return mesh_path.name 

with gr.Blocks() as demo: 
    with gr.Row(variant="panel"): 
        with gr.Column(): 
            with gr.Row(): 
                input_image = gr.Image( 
                    label="Input Image", 
                    image_mode="RGBA", 
                    sources="upload", 
                    type="pil", 
                    elem_id="content_image", 
                ) 
                processed_image = gr.Image(label="Processed Image", interactive=False) 
            with gr.Row(): 
                with gr.Group(): 
                    do_remove_background = gr.Checkbox(label="Remove Background", value=True) 
                    foreground_ratio = gr.Slider( 
                        label="Foreground Ratio", 
                        minimum=0.5, 
                        maximum=1.0, 
                        value=0.85, 
                        step=0.05, 
                    ) 
            with gr.Row(): 
                submit = gr.Button("Generate", elem_id="generate", variant="primary") 
        with gr.Column(): 
            with gr.Tab("Model"): 
                output_model = gr.Model3D( 
                    label="Output Model", 
                    interactive=False, 
                ) 
    with gr.Row(variant="panel"): 
        gr.Examples( 
            examples=[os.path.join("TripoSR/examples", img_name) for img_name in sorted(os.listdir("TripoSR/examples"))], 
            inputs=[input_image], 
            outputs=[processed_image, output_model], 
            label="Examples", 
            examples_per_page=20, 
        ) 
    submit.click(fn=check_input_image, inputs=[input_image]).success( 
        fn=preprocess, 
        inputs=[input_image, do_remove_background, foreground_ratio], 
        outputs=[processed_image], 
    ).success( 
        fn=generate, 
        inputs=[processed_image], 
        outputs=[output_model], 
    ) 

try: 
    demo.launch(debug=False, height=680) 
except Exception: 
    demo.queue().launch(share=True, debug=False, height=680) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細については、ドキュメントをご覧ください: https://gradio.app/docs/