|
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", |
|
"Movement: dollyIn Easing: easeOutBounce Frames: 30 Camera Angle: mediumAngle Shot Type: longShot", |
|
] |
|
DEFAULT_TEXT = [ |
|
"While the character moves right, the camera [...].", |
|
"Movement: dollyIn Easing: easeOutBounce Frames: 30 [...].", |
|
"Movement: shortArcShotRight Easing: easeInOutQuad [...]. " |
|
"Movement: fullZoomIn Easing: easeInOutSine [...].", |
|
] |
|
|
|
HEADER = """ |
|
|
|
<div align="center"> |
|
<h1 style='text-align: center'>Camera Trajectory Generation</h2> |
|
<a href="https://robincourant.github.io/info/"><strong>Robin Courant</strong></a> |
|
路 |
|
<a href="https://nicolas-dufour.github.io/"><strong>Nicolas Dufour</strong></a> |
|
路 |
|
<a href="https://triocrossing.github.io/"><strong>Xi Wang</strong></a> |
|
路 |
|
<a href="http://people.irisa.fr/Marc.Christie/"><strong>Marc Christie</strong></a> |
|
路 |
|
<a href="https://vicky.kalogeiton.info/"><strong>Vicky Kalogeiton</strong></a> |
|
</div> |
|
|
|
|
|
<div align="center"> |
|
<a href="https://www.lix.polytechnique.fr/vista/projects/2024_et_courant/" class="button"><b>[Webpage]</b></a> |
|
<a href="https://github.com/robincourant/DIRECTOR" class="button"><b>[DIRECTOR]</b></a> |
|
<a href="https://github.com/robincourant/CLaTr" class="button"><b>[CLaTr]</b></a> |
|
<a href="https://github.com/robincourant/the-exceptional-trajectories" class="button"><b>[Data]</b></a> |
|
</div> |
|
|
|
<br/> |
|
""" |
|
|
|
|
|
|
|
|
|
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_random_seed(seed) |
|
diffuser.gen_seeds = np.array([seed]) |
|
diffuser.guidance_weight = guidance_weight |
|
|
|
|
|
sample_id = SAMPLE_IDS[0] |
|
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) |
|
|
|
|
|
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, |
|
|
|
) |
|
text = gr.Textbox( |
|
placeholder="Type the camera motion you want to generate", |
|
show_label=True, |
|
label="Text prompt", |
|
value=DEFAULT_TEXT[0], |
|
interactive=True, |
|
|
|
) |
|
seed = gr.Number(value=33, label="Seed") |
|
guidance = gr.Slider(0, 10, value=1.4, label="Guidance", step=0.1) |
|
|
|
|
|
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) |
|
|
