464 lines
17 KiB
Python
464 lines
17 KiB
Python
import argparse, os, sys
|
|
import torch
|
|
import torchvision
|
|
import warnings
|
|
import imageio
|
|
import logging
|
|
import matplotlib.pyplot as plt
|
|
plt.switch_backend('agg')
|
|
import traceback
|
|
import uvicorn
|
|
|
|
from omegaconf import OmegaConf
|
|
from einops import rearrange, repeat
|
|
from collections import OrderedDict
|
|
from pytorch_lightning import seed_everything
|
|
from torch import nn
|
|
from fastapi import FastAPI
|
|
from fastapi.responses import JSONResponse
|
|
from typing import Any, Dict, Optional, Tuple, List
|
|
from datetime import datetime
|
|
|
|
from unifolm_wma.utils.utils import instantiate_from_config
|
|
from unifolm_wma.models.samplers.ddim import DDIMSampler
|
|
|
|
|
|
def get_device_from_parameters(module: nn.Module) -> torch.device:
|
|
"""Get a module's device by checking one of its parameters.
|
|
|
|
Args:
|
|
module (nn.Module): PyTorch module.
|
|
|
|
Returns:
|
|
torch.device: The device where the module's parameters are stored.
|
|
"""
|
|
return next(iter(module.parameters())).device
|
|
|
|
|
|
def load_model_checkpoint(model: nn.Module, ckpt: str) -> nn.Module:
|
|
"""Load model weights from checkpoint file.
|
|
|
|
Args:
|
|
model (nn.Module): Model to load weights into.
|
|
ckpt (str): Path to checkpoint file.
|
|
|
|
Returns:
|
|
nn.Module: Model with loaded weights.
|
|
"""
|
|
|
|
state_dict = torch.load(ckpt, map_location="cpu")
|
|
if "state_dict" in list(state_dict.keys()):
|
|
state_dict = state_dict["state_dict"]
|
|
try:
|
|
model.load_state_dict(state_dict, strict=False)
|
|
except:
|
|
new_pl_sd = OrderedDict()
|
|
for k, v in state_dict.items():
|
|
new_pl_sd[k] = v
|
|
|
|
for k in list(new_pl_sd.keys()):
|
|
if "framestride_embed" in k:
|
|
new_key = k.replace("framestride_embed", "fps_embedding")
|
|
new_pl_sd[new_key] = new_pl_sd[k]
|
|
del new_pl_sd[k]
|
|
model.load_state_dict(new_pl_sd, strict=False)
|
|
else:
|
|
new_pl_sd = OrderedDict()
|
|
for key in state_dict['module'].keys():
|
|
new_pl_sd[key[16:]] = state_dict['module'][key]
|
|
model.load_state_dict(new_pl_sd)
|
|
print('>>> model checkpoint loaded.')
|
|
return model
|
|
|
|
|
|
def write_video(video_path: str, stacked_frames: List[Any], fps: int) -> None:
|
|
"""Write a video to disk using imageio.
|
|
|
|
Args:
|
|
video_path (str): Path to save the video.
|
|
stacked_frames (List[Any]): Frames to write.
|
|
fps (int): Frames per second.
|
|
"""
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore",
|
|
"pkg_resources is deprecated as an API",
|
|
category=DeprecationWarning)
|
|
imageio.mimsave(video_path, stacked_frames, fps=fps)
|
|
|
|
|
|
def save_results(video: torch.Tensor, filename: str, fps: int = 8) -> None:
|
|
"""Save a video tensor as an MP4 file.
|
|
|
|
Args:
|
|
video (torch.Tensor): Video tensor of shape (B, C, T, H, W).
|
|
filename (str): Path to save video.
|
|
fps (int, optional): Frame rate. Defaults to 8.
|
|
|
|
"""
|
|
video = video.detach().cpu()
|
|
video = torch.clamp(video.float(), -1., 1.)
|
|
n = video.shape[0]
|
|
video = video.permute(2, 0, 1, 3, 4)
|
|
|
|
frame_grids = [
|
|
torchvision.utils.make_grid(framesheet, nrow=int(n), padding=0)
|
|
for framesheet in video
|
|
]
|
|
grid = torch.stack(frame_grids, dim=0)
|
|
grid = (grid + 1.0) / 2.0
|
|
grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1)
|
|
torchvision.io.write_video(filename,
|
|
grid,
|
|
fps=fps,
|
|
video_codec='h264',
|
|
options={'crf': '10'})
|
|
|
|
|
|
def get_latent_z(model: nn.Module, videos: torch.Tensor) -> torch.Tensor:
|
|
"""Encode videos into latent space.
|
|
|
|
Args:
|
|
model (nn.Module): Model with `encode_first_stage` method.
|
|
videos (torch.Tensor): Input videos (B, C, T, H, W).
|
|
|
|
Returns:
|
|
torch.Tensor: Latent representation (B, C, T, H, W).
|
|
|
|
"""
|
|
b, c, t, h, w = videos.shape
|
|
x = rearrange(videos, 'b c t h w -> (b t) c h w')
|
|
z = model.encode_first_stage(x)
|
|
z = rearrange(z, '(b t) c h w -> b c t h w', b=b, t=t)
|
|
return z
|
|
|
|
|
|
def image_guided_synthesis(
|
|
model: torch.nn.Module,
|
|
prompts: list[str],
|
|
observation: Dict[str, torch.Tensor],
|
|
noise_shape: tuple[int, int, int, int, int],
|
|
ddim_steps: int = 50,
|
|
ddim_eta: float = 1.0,
|
|
unconditional_guidance_scale: float = 1.0,
|
|
fs: int | None = None,
|
|
timestep_spacing: str = 'uniform',
|
|
guidance_rescale: float = 0.0,
|
|
**kwargs) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Run inference with DDIM sampling.
|
|
|
|
Args:
|
|
model (nn.Module): Diffusion model.
|
|
prompts (Any): Conditioning text prompts.
|
|
observation (Dict[str, torch.Tensor]): Observation dictionary.
|
|
noise_shape (List[int]): Shape of noise tensor.
|
|
ddim_steps (int, optional): Number of DDIM steps. Defaults to 50.
|
|
ddim_eta (float, optional): Sampling eta. Defaults to 1.0.
|
|
unconditional_guidance_scale (float, optional): Guidance scale. Defaults to 1.0.
|
|
fs (Optional[int], optional): Frame stride or FPS. Defaults to None.
|
|
timestep_spacing (str, optional): Spacing strategy. Defaults to "uniform".
|
|
guidance_rescale (float, optional): Guidance rescale. Defaults to 0.0.
|
|
**kwargs (Any): Additional arguments.
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""
|
|
|
|
b, _, t, _, _ = noise_shape
|
|
ddim_sampler = DDIMSampler(model)
|
|
batch_size = noise_shape[0]
|
|
fs = torch.tensor([fs] * batch_size, dtype=torch.long, device=model.device)
|
|
|
|
img = observation['observation.images.top']
|
|
cond_img = img[:, -1, ...]
|
|
cond_img_emb = model.embedder(cond_img)
|
|
cond_img_emb = model.image_proj_model(cond_img_emb)
|
|
|
|
if model.model.conditioning_key == 'hybrid':
|
|
z = get_latent_z(model, img.permute(0, 2, 1, 3, 4))
|
|
img_cat_cond = z[:, :, -1:, :, :]
|
|
img_cat_cond = repeat(img_cat_cond,
|
|
'b c t h w -> b c (repeat t) h w',
|
|
repeat=noise_shape[2])
|
|
cond = {"c_concat": [img_cat_cond]}
|
|
|
|
cond_ins_emb = model.get_learned_conditioning(prompts)
|
|
cond_state = model.state_projector(observation['observation.state'])
|
|
cond_state_emb = model.agent_state_pos_emb + cond_state
|
|
|
|
cond_action = model.action_projector(observation['action'])
|
|
cond_action_emb = model.agent_action_pos_emb + cond_action
|
|
cond_action_emb = torch.zeros_like(cond_action_emb)
|
|
|
|
cond["c_crossattn"] = [
|
|
torch.cat([cond_state_emb, cond_ins_emb, cond_img_emb], dim=1)
|
|
]
|
|
cond["c_crossattn_action"] = [
|
|
observation['observation.images.top'].permute(
|
|
0, 2, 1, 3, 4)[:, :, -model.n_obs_steps_acting:],
|
|
observation['observation.state'][:, -model.n_obs_steps_acting:]
|
|
]
|
|
|
|
uc = None
|
|
|
|
kwargs.update({"unconditional_conditioning_img_nonetext": None})
|
|
|
|
cond_mask = None
|
|
cond_z0 = None
|
|
|
|
if ddim_sampler is not None:
|
|
|
|
samples, actions, states, intermedia = ddim_sampler.sample(
|
|
S=ddim_steps,
|
|
conditioning=cond,
|
|
batch_size=batch_size,
|
|
shape=noise_shape[1:],
|
|
verbose=False,
|
|
unconditional_guidance_scale=unconditional_guidance_scale,
|
|
unconditional_conditioning=uc,
|
|
eta=ddim_eta,
|
|
cfg_img=None,
|
|
mask=cond_mask,
|
|
x0=cond_z0,
|
|
fs=fs,
|
|
timestep_spacing=timestep_spacing,
|
|
guidance_rescale=guidance_rescale,
|
|
**kwargs)
|
|
|
|
# Reconstruct from latent to pixel space
|
|
batch_images = model.decode_first_stage(samples)
|
|
batch_variants = batch_images
|
|
|
|
return batch_variants, actions, states
|
|
|
|
|
|
def run_inference(args: argparse.Namespace, gpu_num: int,
|
|
gpu_no: int) -> Tuple[nn.Module, List[int], Any]:
|
|
"""
|
|
Run inference pipeline on prompts and image inputs.
|
|
|
|
Args:
|
|
args (argparse.Namespace): Parsed command-line arguments.
|
|
gpu_num (int): Number of GPUs.
|
|
gpu_no (int): Index of the current GPU.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
# Load config
|
|
config = OmegaConf.load(args.config)
|
|
# Set use_checkpoint as False as when using deepspeed, it encounters an error "deepspeed backend not set"
|
|
config['model']['params']['wma_config']['params']['use_checkpoint'] = False
|
|
model = instantiate_from_config(config.model)
|
|
model.perframe_ae = args.perframe_ae
|
|
assert os.path.exists(args.ckpt_path), "Error: checkpoint Not Found!"
|
|
model = load_model_checkpoint(model, args.ckpt_path)
|
|
model = model.cuda(gpu_no)
|
|
model.eval()
|
|
print(">>> Model is successfully loaded ...")
|
|
|
|
# Build unnomalizer
|
|
logging.info("***** Configing Data *****")
|
|
data = instantiate_from_config(config.data)
|
|
data.setup()
|
|
print(">>> Dataset is successfully loaded ...")
|
|
|
|
## Run over data
|
|
assert (args.height % 16 == 0) and (
|
|
args.width % 16
|
|
== 0), "Error: image size [h,w] should be multiples of 16!"
|
|
assert args.bs == 1, "Current implementation only support [batch size = 1]!"
|
|
|
|
## Get latent noise shape
|
|
h, w = args.height // 8, args.width // 8
|
|
channels = model.model.diffusion_model.out_channels
|
|
n_frames = args.video_length
|
|
print(f'>>> Generate {n_frames} frames under each generation ...')
|
|
noise_shape = [args.bs, channels, n_frames, h, w]
|
|
|
|
return model, noise_shape, data
|
|
|
|
|
|
def get_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--savedir",
|
|
type=str,
|
|
default=None,
|
|
help="Path to save the results.")
|
|
parser.add_argument("--ckpt_path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to the model checkpoint.")
|
|
parser.add_argument("--config", type=str, help="Path to the config file.")
|
|
parser.add_argument(
|
|
"--ddim_steps",
|
|
type=int,
|
|
default=50,
|
|
help="Number of DDIM steps. If non-positive, DDPM is used instead.")
|
|
parser.add_argument(
|
|
"--ddim_eta",
|
|
type=float,
|
|
default=1.0,
|
|
help="Eta for DDIM sampling. Set to 0.0 for deterministic results.")
|
|
parser.add_argument("--bs",
|
|
type=int,
|
|
default=1,
|
|
help="Batch size for inference. Must be 1.")
|
|
parser.add_argument("--height",
|
|
type=int,
|
|
default=320,
|
|
help="Height of the generated images in pixels.")
|
|
parser.add_argument("--width",
|
|
type=int,
|
|
default=512,
|
|
help="Width of the generated images in pixels.")
|
|
parser.add_argument(
|
|
"--frame_stride",
|
|
type=int,
|
|
default=3,
|
|
help=
|
|
"frame stride control for 256 model (larger->larger motion), FPS control for 512 or 1024 model (smaller->larger motion)"
|
|
)
|
|
parser.add_argument(
|
|
"--unconditional_guidance_scale",
|
|
type=float,
|
|
default=1.0,
|
|
help="Scale for classifier-free guidance during sampling.")
|
|
parser.add_argument("--seed",
|
|
type=int,
|
|
default=123,
|
|
help="Random seed for reproducibility.")
|
|
parser.add_argument("--video_length",
|
|
type=int,
|
|
default=16,
|
|
help="Number of frames in the generated video.")
|
|
parser.add_argument(
|
|
"--timestep_spacing",
|
|
type=str,
|
|
default="uniform",
|
|
help=
|
|
"Strategy for timestep scaling. See Table 2 in the paper: 'Common Diffusion Noise Schedules and Sample Steps are Flawed' (https://huggingface.co/papers/2305.08891)."
|
|
)
|
|
parser.add_argument(
|
|
"--guidance_rescale",
|
|
type=float,
|
|
default=0.0,
|
|
help=
|
|
"Rescale factor for guidance as discussed in 'Common Diffusion Noise Schedules and Sample Steps are Flawed' (https://huggingface.co/papers/2305.08891)."
|
|
)
|
|
parser.add_argument(
|
|
"--perframe_ae",
|
|
action='store_true',
|
|
default=False,
|
|
help=
|
|
"Use per-frame autoencoder decoding to reduce GPU memory usage. Recommended for models with resolutions like 576x1024."
|
|
)
|
|
return parser
|
|
|
|
|
|
class Server:
|
|
|
|
def __init__(self, args: argparse.Namespace) -> None:
|
|
self.model_, self.noise_shape_, self.data_ = run_inference(args, 1, 0)
|
|
self.args_ = args
|
|
self.dataset_name = self.data_.dataset_configs['test']['params'][
|
|
'dataset_name']
|
|
self.device_ = get_device_from_parameters(self.model_)
|
|
|
|
def normalize_image(self, image: torch.Tensor) -> torch.Tensor:
|
|
return (image / 255 - 0.5) * 2
|
|
|
|
def predict_action(self, payload: Dict[str, Any]) -> Any:
|
|
try:
|
|
images = payload['observation.images.top']
|
|
states = payload['observation.state']
|
|
actions = payload['action'] # Should be all zeros
|
|
language_instruction = payload['language_instruction']
|
|
|
|
images = torch.tensor(images).cuda()
|
|
images = self.data_.test_datasets[
|
|
self.dataset_name].spatial_transform(images).unsqueeze(0)
|
|
images = self.normalize_image(images)
|
|
print(f"images shape: {images.shape} ...")
|
|
states = torch.tensor(states)
|
|
states = self.data_.test_datasets[self.dataset_name].normalizer(
|
|
{'observation.state': states})['observation.state']
|
|
states, _ = self.data_.test_datasets[
|
|
self.dataset_name]._map_to_uni_state(states, "joint position")
|
|
print(f"states shape: {states.shape} ...")
|
|
actions = torch.tensor(actions)
|
|
actions, action_mask = self.data_.test_datasets[
|
|
self.dataset_name]._map_to_uni_action(actions,
|
|
"joint position")
|
|
print(f"actions shape: {actions.shape} ...")
|
|
print("=" * 20)
|
|
states = states.unsqueeze(0).cuda()
|
|
actions = actions.unsqueeze(0).cuda()
|
|
|
|
observation = {
|
|
'observation.images.top': images,
|
|
'observation.state': states,
|
|
'action': actions
|
|
}
|
|
observation = {
|
|
key: observation[key].to(self.device_, non_blocking=True)
|
|
for key in observation
|
|
}
|
|
|
|
args = self.args_
|
|
pred_videos, pred_action, _ = image_guided_synthesis(
|
|
self.model_,
|
|
language_instruction,
|
|
observation,
|
|
self.noise_shape_,
|
|
ddim_steps=args.ddim_steps,
|
|
ddim_ets=args.ddim_eta,
|
|
unconditional_guidance_scale=args.unconditional_guidance_scale,
|
|
fs=30 / args.frame_stride,
|
|
timestep_spacing=args.timestep_spacing,
|
|
guidance_rescale=args.guidance_rescale)
|
|
|
|
pred_action = pred_action[..., action_mask[0] == 1.0][0].cpu()
|
|
pred_action = self.data_.test_datasets[
|
|
self.dataset_name].unnormalizer({'action':
|
|
pred_action})['action']
|
|
|
|
os.makedirs(args.savedir, exist_ok=True)
|
|
current_time = datetime.now().strftime("%H:%M:%S")
|
|
video_file = f'{args.savedir}/{current_time}.mp4'
|
|
save_results(pred_videos.cpu(), video_file)
|
|
|
|
response = {
|
|
'result': 'ok',
|
|
'action': pred_action.tolist(),
|
|
'desc': 'success'
|
|
}
|
|
return JSONResponse(response)
|
|
|
|
except:
|
|
logging.error(traceback.format_exc())
|
|
logging.warning(
|
|
"Your request threw an error; make sure your request complies with the expected format:\n"
|
|
"{'image': np.ndarray, 'instruction': str}\n"
|
|
"You can optionally an `unnorm_key: str` to specific the dataset statistics you want to use for "
|
|
"de-normalizing the output actions.")
|
|
return {'result': 'error', 'desc': traceback.format_exc()}
|
|
|
|
def run(self, host: str = "127.0.0.1", port: int = 8000) -> None:
|
|
self.app = FastAPI()
|
|
self.app.post("/predict_action")(self.predict_action)
|
|
print(">>> Inference server is ready ... ")
|
|
uvicorn.run(self.app, host=host, port=port)
|
|
print(">>> Inference server stops ... ")
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = get_parser()
|
|
args = parser.parse_args()
|
|
seed = args.seed
|
|
seed_everything(seed)
|
|
rank, gpu_num = 0, 1
|
|
print(">>> Launch inference server ... ")
|
|
server = Server(args)
|
|
server.run()
|