| import os |
| import sys |
| sys.path.append(os.path.join(os.path.dirname(__file__), '..')) |
| import copy |
| import json |
| import argparse |
| import torch |
| import numpy as np |
| import pandas as pd |
| from tqdm import tqdm |
| from easydict import EasyDict as edict |
| from concurrent.futures import ThreadPoolExecutor |
| from queue import Queue |
|
|
| import trellis.models as models |
| import trellis.modules.sparse as sp |
|
|
|
|
| torch.set_grad_enabled(False) |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--output_dir', type=str, required=True, |
| help='Directory to save the metadata') |
| parser.add_argument('--filter_low_aesthetic_score', type=float, default=None, |
| help='Filter objects with aesthetic score lower than this value') |
| parser.add_argument('--feat_model', type=str, default='dinov2_vitl14_reg', |
| help='Feature model') |
| parser.add_argument('--enc_pretrained', type=str, default='JeffreyXiang/TRELLIS-image-large/ckpts/slat_enc_swin8_B_64l8_fp16', |
| help='Pretrained encoder model') |
| parser.add_argument('--model_root', type=str, default='results', |
| help='Root directory of models') |
| parser.add_argument('--enc_model', type=str, default=None, |
| help='Encoder model. if specified, use this model instead of pretrained model') |
| parser.add_argument('--ckpt', type=str, default=None, |
| help='Checkpoint to load') |
| parser.add_argument('--instances', type=str, default=None, |
| help='Instances to process') |
| parser.add_argument('--rank', type=int, default=0) |
| parser.add_argument('--world_size', type=int, default=1) |
| opt = parser.parse_args() |
| opt = edict(vars(opt)) |
|
|
| if opt.enc_model is None: |
| latent_name = f'{opt.feat_model}_{opt.enc_pretrained.split("/")[-1]}' |
| encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda() |
| else: |
| latent_name = f'{opt.feat_model}_{opt.enc_model}_{opt.ckpt}' |
| cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r'))) |
| encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda() |
| ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt') |
| encoder.load_state_dict(torch.load(ckpt_path), strict=False) |
| encoder.eval() |
| print(f'Loaded model from {ckpt_path}') |
| |
| os.makedirs(os.path.join(opt.output_dir, 'latents', latent_name), exist_ok=True) |
|
|
| |
| if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')): |
| metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv')) |
| else: |
| raise ValueError('metadata.csv not found') |
| if opt.instances is not None: |
| with open(opt.instances, 'r') as f: |
| sha256s = [line.strip() for line in f] |
| metadata = metadata[metadata['sha256'].isin(sha256s)] |
| else: |
| if opt.filter_low_aesthetic_score is not None: |
| metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score] |
| metadata = metadata[metadata[f'feature_{opt.feat_model}'] == True] |
| if f'latent_{latent_name}' in metadata.columns: |
| metadata = metadata[metadata[f'latent_{latent_name}'] == False] |
|
|
| start = len(metadata) * opt.rank // opt.world_size |
| end = len(metadata) * (opt.rank + 1) // opt.world_size |
| metadata = metadata[start:end] |
| records = [] |
| |
| |
| sha256s = list(metadata['sha256'].values) |
| for sha256 in copy.copy(sha256s): |
| if os.path.exists(os.path.join(opt.output_dir, 'latents', latent_name, f'{sha256}.npz')): |
| records.append({'sha256': sha256, f'latent_{latent_name}': True}) |
| sha256s.remove(sha256) |
|
|
| |
| load_queue = Queue(maxsize=4) |
| try: |
| with ThreadPoolExecutor(max_workers=32) as loader_executor, \ |
| ThreadPoolExecutor(max_workers=32) as saver_executor: |
| def loader(sha256): |
| try: |
| feats = np.load(os.path.join(opt.output_dir, 'features', opt.feat_model, f'{sha256}.npz')) |
| load_queue.put((sha256, feats)) |
| except Exception as e: |
| print(f"Error loading features for {sha256}: {e}") |
| loader_executor.map(loader, sha256s) |
| |
| def saver(sha256, pack): |
| save_path = os.path.join(opt.output_dir, 'latents', latent_name, f'{sha256}.npz') |
| np.savez_compressed(save_path, **pack) |
| records.append({'sha256': sha256, f'latent_{latent_name}': True}) |
| |
| for _ in tqdm(range(len(sha256s)), desc="Extracting latents"): |
| sha256, feats = load_queue.get() |
| feats = sp.SparseTensor( |
| feats = torch.from_numpy(feats['patchtokens']).float(), |
| coords = torch.cat([ |
| torch.zeros(feats['patchtokens'].shape[0], 1).int(), |
| torch.from_numpy(feats['indices']).int(), |
| ], dim=1), |
| ).cuda() |
| latent = encoder(feats, sample_posterior=False) |
| assert torch.isfinite(latent.feats).all(), "Non-finite latent" |
| pack = { |
| 'feats': latent.feats.cpu().numpy().astype(np.float32), |
| 'coords': latent.coords[:, 1:].cpu().numpy().astype(np.uint8), |
| } |
| saver_executor.submit(saver, sha256, pack) |
| |
| saver_executor.shutdown(wait=True) |
| except: |
| print("Error happened during processing.") |
| |
| records = pd.DataFrame.from_records(records) |
| records.to_csv(os.path.join(opt.output_dir, f'latent_{latent_name}_{opt.rank}.csv'), index=False) |
|
|
