import spaces
import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
et_dir = os.path.join(current_dir, 'ET')
ccd_dir = os.path.join(current_dir, 'CCD')
sys.path.append(et_dir)
sys.path.append(ccd_dir)
from functools import partial
from typing import Any, Callable, Dict
import clip
import gradio as gr
from gradio_rerun import Rerun
import numpy as np
import trimesh
import rerun as rr
import torch
from ET.utils.common_viz import init, get_batch
from ET.utils.random_utils import set_random_seed
from ET.utils.rerun import et_log_sample
from ET.src.diffuser import Diffuser
from ET.src.datasets.multimodal_dataset import MultimodalDataset
from CCD.utils.rerun import ccd_log_sample
from CCD.src.main import generate_CCD_sample
# ------------------------------------------------------------------------------------- #
batch_size, num_cams, num_verts = None, None, None
SAMPLE_IDS = [
"2011_KAeAqaA0Llg_00005_00001",
"2011_F_EuMeT2wBo_00014_00001",
"2011_MCkKihQrNA4_00014_00000",
]
LABEL_TO_IDS = {
"right": 0,
"static": 1,
"complex": 2,
}
EXAMPLES = [
"While the character moves right, the camera trucks right.",
"While the character moves right, the camera performs a push in.",
"While the character moves right, the camera performs a pull out.",
"The camera pans to the character. The camera switches from right front view to right back view. The character is at the middle center of the screen. The camera shoots at close shot.",
"Movement: fullZoomIn Easing: easeInOutSine Frames: 30 Camera Angle: highAngle Shot Type: closeUp",
"Movement: pedestalDown Easing: easeOutExpo Frames: 30 Camera Angle: mediumAngle Shot Type: longShot", # noqa
"Movement: dollyIn Easing: easeOutBounce Frames: 30 Camera Angle: mediumAngle Shot Type: longShot", # noqa
]
DEFAULT_TEXT = [
"While the character moves right, the camera [...].",
"Movement: dollyIn Easing: easeOutBounce Frames: 30 [...].",
"Movement: shortArcShotRight Easing: easeInOutQuad [...]. "
"Movement: fullZoomIn Easing: easeInOutSine [...].",
]
HEADER = """
"""
# ------------------------------------------------------------------------------------- #
def get_normals(vertices: torch.Tensor, faces: torch.Tensor) -> torch.Tensor:
num_frames, num_faces = vertices.shape[0], faces.shape[-2]
faces = faces.expand(num_frames, num_faces, 3)
normals = [
trimesh.Trimesh(vertices=v, faces=f, process=False).vertex_normals
for v, f in zip(vertices, faces)
]
normals = torch.from_numpy(np.stack(normals))
return normals
@spaces.GPU
def generate_ccd(
prompt: str,
seed: int,
guidance_weight: float,
character_position: list,
) -> Dict[str, Any]:
results = generate_CCD_sample(prompt)
rr.init(f"{3}")
rr.save(".tmp_gr.rrd")
ccd_log_sample(
root_name="world",
traj=np.array(results)
)
return "./.tmp_gr.rrd"
@spaces.GPU
def generate(
prompt: str,
seed: int,
guidance_weight: float,
character_position: list,
# ----------------------- #
dataset: MultimodalDataset,
device: torch.device,
diffuser: Diffuser,
clip_model: clip.model.CLIP,
) -> Dict[str, Any]:
diffuser.to(device)
clip_model.to(device)
# Set arguments
set_random_seed(seed)
diffuser.gen_seeds = np.array([seed])
diffuser.guidance_weight = guidance_weight
# Inference
sample_id = SAMPLE_IDS[0] # Default to the first sample ID
seq_feat = diffuser.net.model.clip_sequential
batch = get_batch(prompt, sample_id, character_position, clip_model, dataset, seq_feat, device)
with torch.no_grad():
out = diffuser.predict_step(batch, 0)
# Run visualization
padding_mask = out["padding_mask"][0].to(bool).cpu()
padded_traj = out["gen_samples"][0].cpu()
traj = padded_traj[padding_mask]
char_traj = out["char_feat"][0].cpu()
padded_vertices = out["char_raw"]["char_vertices"][0]
vertices = padded_vertices[padding_mask]
faces = out["char_raw"]["char_faces"][0]
normals = get_normals(vertices, faces)
fx, fy, cx, cy = out["intrinsics"][0].cpu().numpy()
K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
caption = out["caption_raw"][0]
rr.init(f"{sample_id}")
rr.save(".tmp_gr.rrd")
et_log_sample(
root_name="world",
traj=traj.numpy(),
char_traj=char_traj.numpy(),
K=K,
vertices=vertices.numpy(),
faces=faces.numpy(),
normals=normals.numpy(),
caption=caption,
mesh_masks=None,
)
return "./.tmp_gr.rrd"
# ------------------------------------------------------------------------------------- #
def launch_app(gen_fn_et: Callable, gen_fn_ccd: Callable):
theme = gr.themes.Default(primary_hue="blue", secondary_hue="gray")
model_options = {"ET":gen_fn_et, "CCD":gen_fn_ccd, "LensCraft":gen_fn_et}
with gr.Blocks(theme=theme) as demo:
gr.Markdown(HEADER)
with gr.Row():
with gr.Column(scale=3):
with gr.Column(scale=2):
char_position = gr.Textbox(
placeholder="Enter character position as [x, y, z]",
show_label=True,
label="Character Position (3D vector)",
value="[0.0, 0.0, 0.0]",
interactive=True, # Ensure this is set to True
)
text = gr.Textbox(
placeholder="Type the camera motion you want to generate",
show_label=True,
label="Text prompt",
value=DEFAULT_TEXT[0],
interactive=True, # Ensure this is set to True
)
seed = gr.Number(value=33, label="Seed")
guidance = gr.Slider(0, 10, value=1.4, label="Guidance", step=0.1)
# Add a dropdown menu for selecting the generation model
model_selector = gr.Dropdown(
choices=list(model_options.keys()),
value=list(model_options.keys())[0],
label="Generation Model",
)
with gr.Column(scale=1):
btn = gr.Button("Generate", variant="primary")
with gr.Column(scale=2):
examples = gr.Examples(
examples=[[x, None, None] for x in EXAMPLES],
inputs=[text],
)
with gr.Row():
output = Rerun()
def load_example(example_id):
processed_example = examples.non_none_processed_examples[example_id]
return gr.utils.resolve_singleton(processed_example)
def dynamic_generate(selected_model, *args):
gen_fn = model_options[selected_model]
return gen_fn(*args)
inputs = [text, seed, guidance, char_position]
examples.dataset.click(
load_example,
inputs=[examples.dataset],
outputs=examples.inputs_with_examples,
show_progress=False,
postprocess=False,
queue=False,
).then(fn=dynamic_generate,
inputs=[model_selector, text, seed, guidance, char_position],
outputs=[output])
btn.click(
fn=dynamic_generate,
inputs=[model_selector, text, seed, guidance, char_position],
outputs=[output],
)
text.submit(
fn=dynamic_generate,
inputs=[model_selector, text, seed, guidance, char_position],
outputs=[output],
)
demo.queue().launch(share=False)
# ------------------------------------------------------------------------------------- #
diffuser, clip_model, dataset, device = init("config")
generate_sample_et = partial(
generate,
dataset=dataset,
device=device,
diffuser=diffuser,
clip_model=clip_model,
)
generate_sample_ccd = partial(
generate_ccd
)
launch_app(generate_sample_et, generate_sample_ccd)