app_demo.py

import gradio as gr
import os
import subprocess
import shlex
import spaces
import torch
access_token = os.getenv("HUGGINGFACE_TOKEN")
subprocess.run(
    shlex.split(
        "pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt240/download.html"
    )
)

subprocess.run(
    shlex.split(
        "pip install ./extension/nvdiffrast-0.3.1+torch-py3-none-any.whl --force-reinstall --no-deps"
    )
)

subprocess.run(
    shlex.split(
        "pip install ./extension/renderutils_plugin-0.1.0-cp310-cp310-linux_x86_64.whl --force-reinstall --no-deps"
    )
)
def install_cuda_toolkit():
    # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
    # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])

    os.environ["CUDA_HOME"] = "/usr/local/cuda"
    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
        os.environ["CUDA_HOME"],
        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
    )
    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
    print("==> finfish install")
install_cuda_toolkit()
@spaces.GPU
def check_gpu():
    os.environ['CUDA_HOME'] = '/usr/local/cuda-12.1'
    os.environ['PATH'] += ':/usr/local/cuda-12.1/bin'
    # os.environ['LD_LIBRARY_PATH'] += ':/usr/local/cuda-12.1/lib64'
    os.environ['LD_LIBRARY_PATH'] = "/usr/local/cuda-12.1/lib64:" + os.environ.get('LD_LIBRARY_PATH', '')
    subprocess.run(['nvidia-smi'])  # 测试 CUDA 是否可用
    print(f"torch.cuda.is_available:{torch.cuda.is_available()}")
check_gpu()

from PIL import Image
from einops import rearrange
from diffusers import FluxPipeline
from models.lrm.utils.camera_util import get_flux_input_cameras
from models.lrm.utils.infer_util import save_video
from models.lrm.utils.mesh_util import save_obj, save_obj_with_mtl
from models.lrm.utils.render_utils import rotate_x, rotate_y
from models.lrm.utils.train_util import instantiate_from_config
from models.ISOMER.reconstruction_func import reconstruction
from models.ISOMER.projection_func import projection
import os
from einops import rearrange
from omegaconf import OmegaConf
import torch
import numpy as np
import trimesh
import torchvision
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from torchvision.transforms import v2
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from diffusers import FluxPipeline
from pytorch_lightning import seed_everything
import os
from huggingface_hub import hf_hub_download


from utils.tool import NormalTransfer, get_background, get_render_cameras_video, load_mipmap, render_frames

device_0 = "cuda"
device_1 = "cuda"
resolution = 512
save_dir = "./outputs"
normal_transfer = NormalTransfer()
isomer_azimuths = torch.from_numpy(np.array([0, 90, 180, 270])).float().to(device_1)
isomer_elevations = torch.from_numpy(np.array([5, 5, 5, 5])).float().to(device_1)
isomer_radius = 4.5
isomer_geo_weights = torch.from_numpy(np.array([1, 0.9, 1, 0.9])).float().to(device_1)
isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(device_1)

# model initialization and loading
# flux
# # taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device_0)
# # good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=access_token).to(device_0)
# flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(device=device_0, dtype=torch.bfloat16)
# # flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, vae=taef1, token=access_token).to(device_0)
# flux_lora_ckpt_path = hf_hub_download(repo_id="LTT/xxx-ckpt", filename="rgb_normal_large.safetensors", repo_type="model", token=access_token)
# flux_pipe.load_lora_weights(flux_lora_ckpt_path)
# flux_pipe.to(device=device_0, dtype=torch.bfloat16)
# torch.cuda.empty_cache()
# flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(flux_pipe)


# lrm
config = OmegaConf.load("./models/lrm/config/PRM_inference.yaml")
model_config = config.model_config
infer_config = config.infer_config
model = instantiate_from_config(model_config)
model_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="final_ckpt.ckpt", repo_type="model")
state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
model.load_state_dict(state_dict, strict=True)
model = model.to(device_1)
torch.cuda.empty_cache()
@spaces.GPU
def lrm_reconstructions(image, input_cameras, save_path=None, name="temp", export_texmap=False, if_save_video=False):
    images = image.unsqueeze(0).to(device_1)
    images = v2.functional.resize(images, 512, interpolation=3, antialias=True).clamp(0, 1)
    # breakpoint()
    with torch.no_grad():
        # get triplane
        planes = model.forward_planes(images, input_cameras)
        
        mesh_path_idx = os.path.join(save_path, f'{name}.obj')

        mesh_out = model.extract_mesh(
            planes,
            use_texture_map=export_texmap,
            **infer_config,
        )
        if export_texmap:
            vertices, faces, uvs, mesh_tex_idx, tex_map = mesh_out
            save_obj_with_mtl(
                vertices.data.cpu().numpy(),
                uvs.data.cpu().numpy(),
                faces.data.cpu().numpy(),
                mesh_tex_idx.data.cpu().numpy(),
                tex_map.permute(1, 2, 0).data.cpu().numpy(),
                mesh_path_idx,
            )
        else:
            vertices, faces, vertex_colors = mesh_out
            save_obj(vertices, faces, vertex_colors, mesh_path_idx)
        print(f"Mesh saved to {mesh_path_idx}")

        render_size = 512
        if if_save_video:
            video_path_idx = os.path.join(save_path, f'{name}.mp4')
            render_size = infer_config.render_resolution
            ENV = load_mipmap("models/lrm/env_mipmap/6")
            materials = (0.0,0.9)
            
            all_mv, all_mvp, all_campos = get_render_cameras_video(
                batch_size=1, 
                M=24, 
                radius=4.5, 
                elevation=(90, 60.0),
                is_flexicubes=True,
                fov=30
            )
            
            frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals, alphas = render_frames(
                model, 
                planes, 
                render_cameras=all_mvp,
                camera_pos=all_campos,
                env=ENV,
                materials=materials,
                render_size=render_size, 
                chunk_size=20, 
                is_flexicubes=True,
            )
            normals = (torch.nn.functional.normalize(normals) + 1) / 2
            normals = normals * alphas + (1-alphas)
            all_frames = torch.cat([frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals], dim=3)
                
            save_video(
                all_frames,
                video_path_idx,
                fps=30,
            )
            print(f"Video saved to {video_path_idx}")

    return vertices, faces


def local_normal_global_transform(local_normal_images, azimuths_deg, elevations_deg):
    if local_normal_images.min() >= 0:
        local_normal = local_normal_images.float() * 2 - 1
    else:
        local_normal = local_normal_images.float()
    global_normal = normal_transfer.trans_local_2_global(local_normal, azimuths_deg, elevations_deg, radius=4.5, for_lotus=False)
    global_normal[...,0] *= -1
    global_normal = (global_normal + 1) / 2
    global_normal = global_normal.permute(0, 3, 1, 2)
    return global_normal

# 生成多视图图像
@spaces.GPU(duration=120)
def generate_multi_view_images(prompt, seed):
    # torch.cuda.empty_cache()
    # generator = torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    with torch.no_grad():
        img = flux_pipe(
            prompt=prompt,
            num_inference_steps=5,
            guidance_scale=3.5,
            num_images_per_prompt=1,
            width=resolution * 2,
            height=resolution * 1,
            output_type='np',
            generator=generator,
        ).images
        # for img in flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images(
        #     prompt=prompt,
        #     guidance_scale=3.5,
        #     num_inference_steps=4,
        #     width=resolution * 4,
        #     height=resolution * 2,
        #     generator=generator,
        #     output_type="np",
        #     good_vae=good_vae,
        # ):
        #     pass
    # 返回最终的图像和种子（通过外部调用处理）
    return img

# 重建 3D 模型
@spaces.GPU
def reconstruct_3d_model(images, prompt):
    global model
    model.init_flexicubes_geometry(device_1, fovy=50.0)
    model = model.eval()
    rgb_normal_grid = images
    save_dir_path = os.path.join(save_dir, prompt.replace(" ", "_"))
    os.makedirs(save_dir_path, exist_ok=True)

    images = torch.from_numpy(rgb_normal_grid).squeeze(0).permute(2, 0, 1).contiguous().float()     # (3, 1024, 2048)
    images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=2, m=4)        # (8, 3, 512, 512)
    rgb_multi_view = images[:4, :3, :, :]
    normal_multi_view = images[4:, :3, :, :]
    multi_view_mask = get_background(normal_multi_view)
    rgb_multi_view = rgb_multi_view * rgb_multi_view + (1-multi_view_mask)
    input_cameras = get_flux_input_cameras(batch_size=1, radius=4.2, fov=30).to(device_1)
    vertices, faces = lrm_reconstructions(rgb_multi_view, input_cameras, save_path=save_dir_path, name='lrm', export_texmap=False, if_save_video=True)
    # local normal to global normal

    global_normal = local_normal_global_transform(normal_multi_view.permute(0, 2, 3, 1), isomer_azimuths, isomer_elevations)
    global_normal = global_normal * multi_view_mask + (1-multi_view_mask)

    global_normal = global_normal.permute(0,2,3,1)
    rgb_multi_view = rgb_multi_view.permute(0,2,3,1)
    multi_view_mask = multi_view_mask.permute(0,2,3,1).squeeze(-1)
    vertices = torch.from_numpy(vertices).to(device_1)
    faces = torch.from_numpy(faces).to(device_1)
    vertices = vertices @ rotate_x(np.pi / 2, device=vertices.device)[:3, :3]
    vertices = vertices @ rotate_y(np.pi / 2, device=vertices.device)[:3, :3]

    # global_normal: B,H,W,3
    # multi_view_mask: B,H,W
    # rgb_multi_view: B,H,W,3
    
    meshes = reconstruction(
        normal_pils=global_normal, 
        masks=multi_view_mask, 
        weights=isomer_geo_weights, 
        fov=30, 
        radius=isomer_radius, 
        camera_angles_azi=isomer_azimuths, 
        camera_angles_ele=isomer_elevations, 
        expansion_weight_stage1=0.1,
        init_type="file",
        init_verts=vertices,
        init_faces=faces,
        stage1_steps=0,
        stage2_steps=50,
        start_edge_len_stage1=0.1,
        end_edge_len_stage1=0.02,
        start_edge_len_stage2=0.02,
        end_edge_len_stage2=0.005,
    )


    save_glb_addr = projection(
        meshes,
        masks=multi_view_mask,
        images=rgb_multi_view,
        azimuths=isomer_azimuths, 
        elevations=isomer_elevations, 
        weights=isomer_color_weights,
        fov=30,
        radius=isomer_radius,
        save_dir=f"{save_dir_path}/ISOMER/",
    )

    return save_glb_addr

# Gradio 接口函数
@spaces.GPU
def gradio_pipeline(prompt, seed):
    import ctypes
    # 显式加载 libnvrtc.so.12
    cuda_lib_path = "/usr/local/cuda-12.1/lib64/libnvrtc.so.12"
    try:
        ctypes.CDLL(cuda_lib_path, mode=ctypes.RTLD_GLOBAL)
        print(f"Successfully preloaded {cuda_lib_path}")
    except OSError as e:
        print(f"Failed to preload {cuda_lib_path}: {e}")
    # 生成多视图图像
    # rgb_normal_grid = generate_multi_view_images(prompt, seed)
    rgb_normal_grid = np.load("rgb_normal_grid.npy")
    image_preview = Image.fromarray((rgb_normal_grid[0] * 255).astype(np.uint8))

    # 3d reconstruction


    # 重建 3D 模型并返回 glb 路径
    save_glb_addr = reconstruct_3d_model(rgb_normal_grid, prompt)
    # save_glb_addr = None
    return image_preview, save_glb_addr

# Gradio Blocks 应用
with gr.Blocks() as demo:
    with gr.Row(variant="panel"):
        # 左侧输入区域
        with gr.Column():
            with gr.Row():
                prompt_input = gr.Textbox(
                    label="Enter Prompt",
                    placeholder="Describe your 3D model...",
                    lines=2,
                    elem_id="prompt_input"
                )

            with gr.Row():
                sample_seed = gr.Number(value=42, label="Seed Value", precision=0)

            with gr.Row():
                submit = gr.Button("Generate", elem_id="generate", variant="primary")

            with gr.Row(variant="panel"):
                gr.Markdown("Examples:")
                gr.Examples(
                    examples=[
                        ["a castle on a hill"],
                        ["an owl wearing a hat"],
                        ["a futuristic car"]
                    ],
                    inputs=[prompt_input],
                    label="Prompt Examples"
                )

        # 右侧输出区域
        with gr.Column():
            with gr.Row():
                rgb_normal_grid_image = gr.Image(
                    label="RGB Normal Grid",
                    type="pil",
                    interactive=False
                )

            with gr.Row():
                with gr.Tab("GLB"):
                    output_glb_model = gr.Model3D(
                        label="Generated 3D Model (GLB Format)",
                        interactive=False
                    )
                    gr.Markdown("Download the model for proper visualization.")

    # 处理逻辑
    submit.click(
        fn=gradio_pipeline, inputs=[prompt_input, sample_seed],
        outputs=[rgb_normal_grid_image, output_glb_model]
    )

# 启动应用
# demo.queue(max_size=10)
demo.launch()