File size: 5,280 Bytes
6b92ff7 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import os
import torch
import argparse
from PIL import Image
# Add current directory to path to allow imports
import sys
sys.path.append(os.getcwd())
sys.path.append(os.path.join(os.getcwd(), 'third_parties/dsine'))
from anigen.pipelines import AnigenImageTo3DPipeline
from anigen.utils.random_utils import set_random_seed
from anigen.utils.image_utils import _expand_image_inputs
from anigen.utils.ckpt_utils import ensure_ckpts
@torch.no_grad()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, required=True, help='Path to input image or a folder of images')
parser.add_argument('--ss_flow_path', type=str, required=False, default='ckpts/anigen/ss_flow_duet', help='Path to SS Flow model directory')
parser.add_argument('--slat_flow_path', type=str, required=False, default='ckpts/anigen/slat_flow_auto', help='Path to SLat Flow model directory')
parser.add_argument('--output_dir', type=str, default='results/', help='Output directory')
parser.add_argument('--seed', type=int, default=42)
parser.add_argument('--cfg_scale_ss', type=float, default=7.5, help='Classifier-free guidance scale')
parser.add_argument('--cfg_scale', type=float, default=3.0, help='Classifier-free guidance scale')
parser.add_argument('--deterministic', action='store_true', help='Enable mostly-deterministic torch behavior (may be slower)')
parser.add_argument('--device', type=str, default='cuda')
parser.add_argument('--use_ema', action='store_true', help='Use EMA checkpoint if available')
parser.add_argument(
'--output_name',
type=str,
default=None,
help='Optional subfolder name to save outputs under `--output_dir`. If not provided, the image filename stem is used.',
)
parser.add_argument('--no_smooth_skin_weights', action='store_true', help='Disable skin-weight smoothing')
parser.add_argument('--smooth_skin_weights_iters', type=int, default=100, help='Number of smoothing iterations (default: 100)')
parser.add_argument('--smooth_skin_weights_alpha', type=float, default=1.0, help='Smoothing alpha (default: 1.0)')
parser.add_argument(
'--no_filter_skin_weights',
action='store_true',
help='Use geodesic distribution to filter mesh skinning weights.',
)
parser.add_argument(
'--joints_density', '--joint_density',
type=int,
default=1,
help='Optional joint density level for Slat flow (from 0 to 4, higher means more joints)',
)
args = parser.parse_args()
base_output_dir = args.output_dir
input_image_paths, is_dir = _expand_image_inputs(args.image_path)
if is_dir and len(input_image_paths) == 0:
raise ValueError(f"No supported images found under directory: {args.image_path}")
# For directory input, group outputs under a batch folder.
# For single-image input, keep original behavior: output under `<base_output_dir>/<output_name|image_stem>`.
batch_folder_name = None
if is_dir:
batch_folder_name = args.output_name if (args.output_name is not None and str(args.output_name).strip() != '') else os.path.basename(os.path.normpath(args.image_path))
set_random_seed(args.seed, deterministic=args.deterministic)
ensure_ckpts()
print("Loading models...")
pipeline = AnigenImageTo3DPipeline.from_pretrained(
ss_flow_path=args.ss_flow_path,
slat_flow_path=args.slat_flow_path,
device=args.device,
use_ema=args.use_ema
)
pipeline.cuda()
for idx, cur_image_path in enumerate(input_image_paths):
# Per-image output directory.
image_stem = os.path.splitext(os.path.basename(cur_image_path))[0]
if is_dir:
args.output_dir = os.path.join(base_output_dir, str(batch_folder_name), image_stem)
else:
# Allow user to override the saved folder name via --output_name. Fallback to image stem.
folder_name = args.output_name if (args.output_name is not None and str(args.output_name).strip() != '') else image_stem
args.output_dir = os.path.join(base_output_dir, folder_name)
os.makedirs(args.output_dir, exist_ok=True)
# Keep args.image_path aligned for any downstream logging/debug usage.
args.image_path = cur_image_path
print(f"Processing image {idx + 1}/{len(input_image_paths)}: {cur_image_path}")
image = Image.open(cur_image_path)
# Run pipeline
output_glb_path = os.path.join(args.output_dir, 'mesh.glb')
outputs = pipeline.run(
image,
seed=args.seed,
cfg_scale_ss=args.cfg_scale_ss,
cfg_scale_slat=args.cfg_scale,
joints_density=args.joints_density,
no_smooth_skin_weights=args.no_smooth_skin_weights,
no_filter_skin_weights=args.no_filter_skin_weights,
smooth_skin_weights_iters=args.smooth_skin_weights_iters,
smooth_skin_weights_alpha=args.smooth_skin_weights_alpha,
output_glb=output_glb_path
)
# Save processed images
outputs['processed_image'].save(os.path.join(args.output_dir, 'processed_image.png'))
if __name__ == '__main__':
main()
|