Spaces:
Paused
Paused
File size: 2,298 Bytes
02afb14 |
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 |
import argparse
from pathlib import Path
import numpy as np
import torch
from omegaconf import OmegaConf
from skimage.io import imsave
from ldm.models.diffusion.sync_dreamer import SyncMultiviewDiffusion
from ldm.util import instantiate_from_config, prepare_inputs
def load_model(cfg,ckpt,strict=True):
config = OmegaConf.load(cfg)
model = instantiate_from_config(config.model)
print(f'loading model from {ckpt} ...')
ckpt = torch.load(ckpt,map_location='cpu')
model.load_state_dict(ckpt['state_dict'],strict=strict)
model = model.cuda().eval()
return model
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--cfg',type=str, default='configs/syncdreamer.yaml')
parser.add_argument('--ckpt',type=str, default='ckpt/syncdreamer-step80k.ckpt')
parser.add_argument('--output', type=str, required=True)
parser.add_argument('--input', type=str, required=True)
parser.add_argument('--elevation', type=float, required=True)
parser.add_argument('--sample_num', type=int, default=4)
parser.add_argument('--crop_size', type=int, default=-1)
parser.add_argument('--cfg_scale', type=float, default=2.0)
parser.add_argument('--batch_view_num', type=int, default=8)
parser.add_argument('--seed', type=int, default=6033)
flags = parser.parse_args()
torch.random.manual_seed(flags.seed)
np.random.seed(flags.seed)
model = load_model(flags.cfg, flags.ckpt, strict=True)
assert isinstance(model, SyncMultiviewDiffusion)
Path(f'{flags.output}').mkdir(exist_ok=True, parents=True)
# prepare data
data = prepare_inputs(flags.input, flags.elevation, flags.crop_size)
for k, v in data.items():
data[k] = v.unsqueeze(0).cuda()
data[k] = torch.repeat_interleave(data[k], flags.sample_num, dim=0)
x_sample = model.sample(data, flags.cfg_scale, flags.batch_view_num)
B, N, _, H, W = x_sample.shape
x_sample = (torch.clamp(x_sample,max=1.0,min=-1.0) + 1) * 0.5
x_sample = x_sample.permute(0,1,3,4,2).cpu().numpy() * 255
x_sample = x_sample.astype(np.uint8)
for bi in range(B):
output_fn = Path(flags.output)/ f'{bi}.png'
imsave(output_fn, np.concatenate([x_sample[bi,ni] for ni in range(N)], 1))
if __name__=="__main__":
main()
|