File size: 3,237 Bytes
cd7282c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import argparse
import os
import sys
import torch
from mvadapter.pipelines.pipeline_texture import ModProcessConfig, TexturePipeline
from mvadapter.utils import make_image_grid
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--device", type=str, default="cuda")
parser.add_argument("--variant", type=str, default="sdxl", choices=["sdxl", "sd21"])
# I/O
parser.add_argument("--mesh", type=str, required=True)
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--save_dir", type=str, default="./output")
parser.add_argument("--save_name", type=str, default="t2tex_sample")
# Extra
parser.add_argument("--preprocess_mesh", action="store_true")
args = parser.parse_args()
if args.variant == "sdxl":
from .inference_tg2mv_sdxl import prepare_pipeline, run_pipeline
base_model = "stabilityai/stable-diffusion-xl-base-1.0"
vae_model = "madebyollin/sdxl-vae-fp16-fix"
height = width = 768
uv_size = 4096
elif args.variant == "sd21":
from .inference_tg2mv_sd import prepare_pipeline, run_pipeline
base_model = "stabilityai/stable-diffusion-2-1-base"
vae_model = None
height = width = 512
uv_size = 2048
else:
raise ValueError(f"Invalid variant: {args.variant}")
device = args.device
num_views = 6
# Prepare pipelines
pipe = prepare_pipeline(
base_model=base_model,
vae_model=vae_model,
unet_model=None,
lora_model=None,
adapter_path="huanngzh/mv-adapter",
scheduler=None,
num_views=num_views,
device=device,
dtype=torch.float16,
)
texture_pipe = TexturePipeline(
upscaler_ckpt_path="./checkpoints/RealESRGAN_x2plus.pth",
inpaint_ckpt_path="./checkpoints/big-lama.pt",
device=device,
)
print("Pipeline ready.")
os.makedirs(args.save_dir, exist_ok=True)
# 1. run MV-Adapter to generate multi-view images
images, pos_images, normal_images = run_pipeline(
pipe,
mesh_path=args.mesh,
num_views=num_views,
text=args.text,
height=height,
width=width,
num_inference_steps=50,
guidance_scale=7.0,
seed=args.seed,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
device=device,
)
mv_path = os.path.join(args.save_dir, f"{args.save_name}.png")
make_image_grid(images, rows=1).save(mv_path)
torch.cuda.empty_cache()
# 2. un-project and complete texture
out = texture_pipe(
mesh_path=args.mesh,
save_dir=args.save_dir,
save_name=args.save_name,
uv_unwarp=True,
preprocess_mesh=args.preprocess_mesh,
uv_size=uv_size,
rgb_path=mv_path,
rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
)
print(f"Output saved to {out.shaded_model_save_path}")
|