wehub-resource-sync
313 lines
13 KiB
Python
313 lines
13 KiB
Python
import os
|
|
import argparse
|
|
import math
|
|
import time
|
|
import torch
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
|
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
|
|
os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
|
|
os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
|
|
|
|
from pixal3d.pipelines import Pixal3DImageTo3DPipeline
|
|
import o_voxel
|
|
|
|
# ============================================================================
|
|
# Constants & Defaults
|
|
# ============================================================================
|
|
|
|
MOGE_MODEL_NAME = "Ruicheng/moge-2-vitl"
|
|
MODEL_PATH = "TencentARC/Pixal3D"
|
|
|
|
IMAGE_COND_CONFIGS = {
|
|
"ss": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 512,
|
|
"grid_resolution": 16,
|
|
},
|
|
"shape_512": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 512,
|
|
"grid_resolution": 32,
|
|
"use_naf_upsample": True,
|
|
"naf_target_size": 512,
|
|
},
|
|
"shape_1024": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 1024,
|
|
"grid_resolution": 64,
|
|
"use_naf_upsample": True,
|
|
"naf_target_size": 512,
|
|
},
|
|
"tex_1024": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 1024,
|
|
"grid_resolution": 64,
|
|
"use_naf_upsample": True,
|
|
"naf_target_size": 1024,
|
|
},
|
|
}
|
|
|
|
# ============================================================================
|
|
# Model Loading
|
|
# ============================================================================
|
|
|
|
def build_image_cond_model(config: dict):
|
|
from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor
|
|
model = DinoV3ProjFeatureExtractor(**config)
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME):
|
|
from moge.model.v2 import MoGeModel
|
|
moge_model = MoGeModel.from_pretrained(model_name)
|
|
moge_model = moge_model.to(device)
|
|
moge_model.eval()
|
|
return moge_model
|
|
|
|
|
|
def init_pipeline(model_path=MODEL_PATH, device="cuda", low_vram=False):
|
|
print(f"[Pipeline] Loading from {model_path}...")
|
|
pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
|
|
|
|
print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
|
|
pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
|
|
pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
|
|
pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
|
|
pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
|
|
|
|
if low_vram:
|
|
# Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage.
|
|
# Peak VRAM = one flow model + one DinoV3, not all ~18 GB at once.
|
|
print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...")
|
|
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
|
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
|
m = getattr(pipeline, attr, None)
|
|
if m is not None and getattr(m, 'use_naf_upsample', False):
|
|
m._load_naf()
|
|
pipeline._device = torch.device(device)
|
|
pipeline.low_vram = True
|
|
print("[Pipeline] Low-VRAM mode enabled.")
|
|
else:
|
|
# Standard mode: all models loaded to GPU at once (faster, needs more VRAM).
|
|
pipeline.low_vram = False
|
|
pipeline.cuda()
|
|
pipeline.image_cond_model_ss.cuda()
|
|
pipeline.image_cond_model_shape_512.cuda()
|
|
pipeline.image_cond_model_shape_1024.cuda()
|
|
pipeline.image_cond_model_tex_1024.cuda()
|
|
print("[NAF] Pre-loading NAF upsampler model...")
|
|
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
|
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
|
m = getattr(pipeline, attr, None)
|
|
if m is not None and getattr(m, 'use_naf_upsample', False):
|
|
m._load_naf()
|
|
print("[Pipeline] Standard mode (all models on GPU).")
|
|
|
|
return pipeline
|
|
|
|
# ============================================================================
|
|
# Camera Estimation
|
|
# ============================================================================
|
|
|
|
def compute_f_pixels(camera_angle_x: float, resolution: int) -> float:
|
|
focal_length = 16.0 / torch.tan(torch.tensor(camera_angle_x / 2.0))
|
|
f_pixels = focal_length * resolution / 32.0
|
|
return float(f_pixels.item())
|
|
|
|
|
|
def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, image_resolution):
|
|
rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
|
|
gp = grid_point.to(torch.float32) @ rotation_matrix.T
|
|
gp = gp / mesh_scale / 2
|
|
xw, yw, zw = gp[0].item(), gp[1].item(), gp[2].item()
|
|
xt, yt = float(target_point[0].item()), float(target_point[1].item())
|
|
f_pixels = compute_f_pixels(camera_angle_x, image_resolution)
|
|
x_ndc = xt - image_resolution / 2.0
|
|
y_ndc = -(yt - image_resolution / 2.0)
|
|
distance_x = f_pixels * xw / x_ndc - yw
|
|
return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)}
|
|
|
|
|
|
def get_camera_params_wild_moge(image_path, moge_model, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512):
|
|
pil_image = Image.open(image_path).convert("RGB")
|
|
width, height = pil_image.size
|
|
image_np = np.array(pil_image).astype(np.float32) / 255.0
|
|
image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).to(device)
|
|
with torch.no_grad():
|
|
output = moge_model.infer(image_tensor)
|
|
intrinsics = output["intrinsics"].squeeze().cpu().numpy()
|
|
fx_normalized = intrinsics[0, 0]
|
|
fx = fx_normalized * width
|
|
camera_angle_x = 2 * math.atan(width / (2 * fx))
|
|
|
|
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
|
distance = distance_from_fov(
|
|
camera_angle_x, grid_point,
|
|
torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
|
|
mesh_scale, image_resolution
|
|
)["distance_from_x"]
|
|
return {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
|
|
|
|
# ============================================================================
|
|
# Main Inference
|
|
# ============================================================================
|
|
|
|
def run_inference(
|
|
image_path: str,
|
|
output_path: str,
|
|
seed: int = 42,
|
|
ss_guidance_strength: float = 7.5,
|
|
ss_guidance_rescale: float = 0.7,
|
|
ss_sampling_steps: int = 12,
|
|
ss_rescale_t: float = 5.0,
|
|
shape_slat_guidance_strength: float = 7.5,
|
|
shape_slat_guidance_rescale: float = 0.5,
|
|
shape_slat_sampling_steps: int = 12,
|
|
shape_slat_rescale_t: float = 3.0,
|
|
tex_slat_guidance_strength: float = 1.0,
|
|
tex_slat_guidance_rescale: float = 0.0,
|
|
tex_slat_sampling_steps: int = 12,
|
|
tex_slat_rescale_t: float = 3.0,
|
|
mesh_scale: float = 1.0,
|
|
extend_pixel: int = 0,
|
|
image_resolution: int = 512,
|
|
max_num_tokens: int = 49152,
|
|
model_path: str = MODEL_PATH,
|
|
manual_fov: float = -1.0,
|
|
low_vram: bool = False,
|
|
resolution: int = -1,
|
|
):
|
|
# Load models
|
|
pipeline = init_pipeline(model_path, low_vram=low_vram)
|
|
|
|
# Preprocess image first — rembg loads to GPU for this call, then offloads.
|
|
# MoGe is loaded afterwards so both never occupy VRAM at the same time.
|
|
print(f"[Inference] Processing image: {image_path}")
|
|
img = Image.open(image_path)
|
|
image_preprocessed = pipeline.preprocess_image(img)
|
|
|
|
# Save preprocessed image for MoGe
|
|
tmp_path = os.path.join(os.path.dirname(os.path.abspath(output_path)), f"_tmp_preprocessed_{int(time.time()*1000)}.png")
|
|
image_preprocessed.save(tmp_path)
|
|
|
|
# Camera estimation
|
|
if manual_fov > 0:
|
|
# Use manually specified FOV (in radians)
|
|
camera_angle_x = float(manual_fov)
|
|
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
|
distance = distance_from_fov(
|
|
camera_angle_x, grid_point,
|
|
torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
|
|
mesh_scale, image_resolution
|
|
)["distance_from_x"]
|
|
camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
|
|
print(f"[Inference] Using manual FOV: {math.degrees(manual_fov):.2f}° ({manual_fov:.4f} rad), distance={distance:.4f}")
|
|
else:
|
|
print("[MoGe-2] Loading model for camera estimation...")
|
|
moge_model = load_moge_model(device="cuda")
|
|
print("[Inference] Estimating camera parameters...")
|
|
camera_params = get_camera_params_wild_moge(
|
|
tmp_path, moge_model, device="cuda",
|
|
mesh_scale=mesh_scale, extend_pixel=extend_pixel,
|
|
image_resolution=image_resolution,
|
|
)
|
|
print(f" camera_angle_x={camera_params['camera_angle_x']:.4f}, distance={camera_params['distance']:.4f}")
|
|
# MoGe is only needed for camera estimation; free its VRAM for inference.
|
|
moge_model.cpu()
|
|
del moge_model
|
|
torch.cuda.empty_cache()
|
|
os.remove(tmp_path)
|
|
|
|
# Run pipeline
|
|
print("[Inference] Running 3D generation pipeline...")
|
|
torch.manual_seed(seed)
|
|
|
|
ss_sampler_override = {
|
|
"steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength,
|
|
"guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t,
|
|
}
|
|
shape_sampler_override = {
|
|
"steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength,
|
|
"guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t,
|
|
}
|
|
tex_sampler_override = {
|
|
"steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength,
|
|
"guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t,
|
|
}
|
|
|
|
pipeline_type = f"{resolution if resolution > 0 else (1024 if low_vram else 1536)}_cascade"
|
|
print(f"[Inference] Using pipeline_type={pipeline_type}")
|
|
mesh_list, (shape_slat, tex_slat, res) = pipeline.run(
|
|
image_preprocessed,
|
|
camera_params=camera_params,
|
|
seed=seed,
|
|
sparse_structure_sampler_params=ss_sampler_override,
|
|
shape_slat_sampler_params=shape_sampler_override,
|
|
tex_slat_sampler_params=tex_sampler_override,
|
|
preprocess_image=False,
|
|
return_latent=True,
|
|
pipeline_type=pipeline_type,
|
|
max_num_tokens=max_num_tokens,
|
|
)
|
|
|
|
mesh = mesh_list[0]
|
|
|
|
# Extract GLB
|
|
print("[Inference] Extracting GLB...")
|
|
glb = o_voxel.postprocess.to_glb(
|
|
vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs,
|
|
coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout,
|
|
grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
decimation_target=1000000, texture_size=4096,
|
|
remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
|
|
)
|
|
|
|
# Apply rotation
|
|
rot = np.array([
|
|
[-1, 0, 0, 0],
|
|
[ 0, 0, -1, 0],
|
|
[ 0, -1, 0, 0],
|
|
[ 0, 0, 0, 1],
|
|
], dtype=np.float64)
|
|
glb.apply_transform(rot)
|
|
|
|
# Export
|
|
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
|
|
glb.export(output_path, extension_webp=True)
|
|
print(f"[Done] GLB saved to: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Pixal3D Inference: Image to GLB")
|
|
parser.add_argument("--image", type=str, required=True, help="Path to input image")
|
|
parser.add_argument("--output", type=str, default="./output.glb", help="Output GLB file path")
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
|
parser.add_argument("--fov", type=float, default=-1.0,
|
|
help="Manual camera FOV in radians (e.g. 0.2). "
|
|
"If not set, FOV is auto-estimated via MoGe-2. "
|
|
"Try 0.2 rad if you notice distortion.")
|
|
parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Model path or HuggingFace repo")
|
|
parser.add_argument("--low_vram", action="store_true",
|
|
help="Enable low-VRAM mode: models stay on CPU and are loaded to GPU on-demand per stage. "
|
|
"Reduces peak VRAM from ~18GB to ~10-12GB at the cost of slower inference.")
|
|
parser.add_argument("--resolution", type=int, default=-1,
|
|
help="Pipeline resolution (1024 or 1536). Default: 1024 if --low_vram, else 1536.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
run_inference(
|
|
image_path=args.image,
|
|
output_path=args.output,
|
|
seed=args.seed,
|
|
manual_fov=args.fov,
|
|
model_path=args.model_path,
|
|
low_vram=args.low_vram,
|
|
resolution=args.resolution,
|
|
)
|