File size: 3,119 Bytes
040f4b2 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import os
import torch
import torch.nn as nn
import numpy as np
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Dict, Any
app = FastAPI()
# A "Realistic Mock" of the AnyTop Architecture
# Based on the SIGGRAPH 2025 paper "AnyTop: Character Animation Diffusion with Any Topology"
# Features: Enrichment Block, Skeletal Attention, Temporal Attention
class AnyTopMockModel(nn.Module):
def __init__(self, feat_dim=6, latent_dim=128):
super().__init__()
# feat_dim: 3 for pos, 3 for rot (simplified)
self.enrichment = nn.Linear(feat_dim + 3, latent_dim) # +3 for rest-pose embedding
self.skeletal_transformer = nn.TransformerEncoderLayer(d_model=latent_dim, nhead=4, batch_first=True)
self.temporal_transformer = nn.TransformerEncoderLayer(d_model=latent_dim, nhead=4, batch_first=True)
self.output_head = nn.Linear(latent_dim, feat_dim)
def forward(self, motion_noise, rest_pose, topology_dist):
# motion_noise: (batch, frames, joints, feat_dim)
# rest_pose: (batch, joints, 3)
b, f, j, d = motion_noise.shape
# Enrichment: Concat rest pose to features
# Broadcast rest_pose to all frames
rp_expanded = rest_pose.unsqueeze(1).expand(-1, f, -1, -1)
x = torch.cat([motion_noise, rp_expanded], dim=-1)
x = self.enrichment(x) # (b, f, j, latent_dim)
# Skeletal Attention (per frame)
x = x.view(b * f, j, -1)
# In a real model, topology_dist would bias the attention matrix here
x = self.skeletal_transformer(x)
x = x.view(b, f, j, -1)
# Temporal Attention (per joint)
x = x.transpose(1, 2).reshape(b * j, f, -1)
x = self.temporal_transformer(x)
x = x.view(b, j, f, -1).transpose(1, 2)
return self.output_head(x)
# Global Instance
model = AnyTopMockModel()
class InferenceRequest(BaseModel):
skeleton: Dict[str, Any]
steps: int = 10
@app.post("/infer")
async def infer(data: InferenceRequest):
# Simulate Diffusion Iteration
# Input: joints and hierarchy
joints = data.skeleton.get("joints", [])
num_joints = len(joints)
num_frames = 60
# Random Noise Start
motion = torch.randn(1, num_frames, num_joints, 6)
rest_pose = torch.tensor([j["pos"] for j in joints]).unsqueeze(0).float()
# Mocking the denoising loop
with torch.no_grad():
for _ in range(min(data.steps, 5)): # Max 5 for speed in demo
motion = model(motion, rest_pose, None)
# Convert to list for JSON response
# Shape: (frames, joints, XYZ)
result = motion.squeeze(0).numpy()[..., :3].tolist()
return {"motion": result}
app.mount("/static", StaticFiles(directory="static"), name="static")
@app.get("/", response_class=HTMLResponse)
async def index():
with open("templates/index.html") as f:
return f.read()
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|