565 lines
17 KiB
Rust
565 lines
17 KiB
Rust
//! Sample data for WorldGen demo.
|
|
|
|
use worldgen_shared::{
|
|
GenerationRequest, LightingState, ModelConfig, ModelType, PhysicsConstraint,
|
|
PhysicsConstraintType, PhysicsParameters, SchedulerConfig, SchedulerType, TrainingConfig,
|
|
WorldObject, WorldState,
|
|
};
|
|
|
|
/// Create a sample generation request for a bouncing ball.
|
|
#[must_use]
|
|
pub fn bouncing_ball_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "A red ball bouncing on a wooden floor in a sunlit room".to_string(),
|
|
negative_prompt: Some("blurry, distorted, unrealistic physics".to_string()),
|
|
num_frames: 24,
|
|
width: 256,
|
|
height: 256,
|
|
fps: 12.0,
|
|
physics_constraints: vec![
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Gravity,
|
|
weight: 1.0,
|
|
parameters: PhysicsParameters {
|
|
gravity: Some([0.0, -9.81, 0.0]),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Collision,
|
|
weight: 0.9,
|
|
parameters: PhysicsParameters {
|
|
restitution: Some(0.8),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::TemporalCoherence,
|
|
weight: 0.7,
|
|
parameters: Default::default(),
|
|
},
|
|
],
|
|
seed: Some(42),
|
|
guidance_scale: 7.5,
|
|
num_inference_steps: 50,
|
|
}
|
|
}
|
|
|
|
/// Create a sample generation request for flowing water.
|
|
#[must_use]
|
|
pub fn flowing_water_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "Crystal clear water flowing over smooth river rocks, sunlight sparkling"
|
|
.to_string(),
|
|
negative_prompt: Some("frozen, static, unrealistic".to_string()),
|
|
num_frames: 32,
|
|
width: 512,
|
|
height: 288,
|
|
fps: 24.0,
|
|
physics_constraints: vec![
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Fluid,
|
|
weight: 1.0,
|
|
parameters: PhysicsParameters {
|
|
viscosity: Some(0.001), // Water viscosity
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Lighting,
|
|
weight: 0.8,
|
|
parameters: PhysicsParameters {
|
|
light_direction: Some([0.3, -0.8, 0.5]),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::TemporalCoherence,
|
|
weight: 0.9,
|
|
parameters: Default::default(),
|
|
},
|
|
],
|
|
seed: Some(123),
|
|
guidance_scale: 8.0,
|
|
num_inference_steps: 75,
|
|
}
|
|
}
|
|
|
|
/// Create a sample generation request for a pendulum.
|
|
#[must_use]
|
|
pub fn pendulum_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "A brass pendulum swinging in a grandfather clock, smooth motion".to_string(),
|
|
negative_prompt: Some("jerky, unnatural motion".to_string()),
|
|
num_frames: 48,
|
|
width: 256,
|
|
height: 384,
|
|
fps: 30.0,
|
|
physics_constraints: vec![
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Gravity,
|
|
weight: 1.0,
|
|
parameters: PhysicsParameters {
|
|
gravity: Some([0.0, -9.81, 0.0]),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::RigidBody,
|
|
weight: 1.0,
|
|
parameters: Default::default(),
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Momentum,
|
|
weight: 0.9,
|
|
parameters: Default::default(),
|
|
},
|
|
],
|
|
seed: Some(456),
|
|
guidance_scale: 7.0,
|
|
num_inference_steps: 60,
|
|
}
|
|
}
|
|
|
|
/// Create a sample generation request for a car driving.
|
|
#[must_use]
|
|
pub fn driving_car_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "A silver sports car driving along a coastal highway at sunset".to_string(),
|
|
negative_prompt: Some("floating, unrealistic shadows, glitchy".to_string()),
|
|
num_frames: 64,
|
|
width: 512,
|
|
height: 288,
|
|
fps: 24.0,
|
|
physics_constraints: vec![
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::RigidBody,
|
|
weight: 1.0,
|
|
parameters: Default::default(),
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Shadows,
|
|
weight: 0.8,
|
|
parameters: Default::default(),
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Lighting,
|
|
weight: 0.9,
|
|
parameters: PhysicsParameters {
|
|
light_direction: Some([-0.5, -0.3, 0.8]),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::TemporalCoherence,
|
|
weight: 0.95,
|
|
parameters: Default::default(),
|
|
},
|
|
],
|
|
seed: Some(789),
|
|
guidance_scale: 8.5,
|
|
num_inference_steps: 80,
|
|
}
|
|
}
|
|
|
|
/// Create a sample generation request for a jellyfish.
|
|
#[must_use]
|
|
pub fn jellyfish_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "A bioluminescent jellyfish pulsing through dark ocean water".to_string(),
|
|
negative_prompt: Some("static, rigid, unrealistic".to_string()),
|
|
num_frames: 40,
|
|
width: 384,
|
|
height: 384,
|
|
fps: 20.0,
|
|
physics_constraints: vec![
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::SoftBody,
|
|
weight: 1.0,
|
|
parameters: Default::default(),
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Fluid,
|
|
weight: 0.7,
|
|
parameters: PhysicsParameters {
|
|
viscosity: Some(0.01),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Lighting,
|
|
weight: 0.6,
|
|
parameters: Default::default(),
|
|
},
|
|
],
|
|
seed: Some(999),
|
|
guidance_scale: 9.0,
|
|
num_inference_steps: 70,
|
|
}
|
|
}
|
|
|
|
/// Create a sample generation request for falling leaves.
|
|
#[must_use]
|
|
pub fn falling_leaves_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "Autumn leaves gently falling from a maple tree in golden afternoon light"
|
|
.to_string(),
|
|
negative_prompt: Some("unnatural motion, frozen".to_string()),
|
|
num_frames: 60,
|
|
width: 384,
|
|
height: 512,
|
|
fps: 24.0,
|
|
physics_constraints: vec![
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Gravity,
|
|
weight: 0.6,
|
|
parameters: PhysicsParameters {
|
|
gravity: Some([0.0, -2.0, 0.0]), // Slower falling
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Fluid,
|
|
weight: 0.4, // Air resistance
|
|
parameters: PhysicsParameters {
|
|
viscosity: Some(0.1),
|
|
..Default::default()
|
|
},
|
|
},
|
|
PhysicsConstraint {
|
|
constraint_type: PhysicsConstraintType::Lighting,
|
|
weight: 0.8,
|
|
parameters: PhysicsParameters {
|
|
light_direction: Some([0.4, -0.6, 0.7]),
|
|
..Default::default()
|
|
},
|
|
},
|
|
],
|
|
seed: Some(1111),
|
|
guidance_scale: 7.5,
|
|
num_inference_steps: 55,
|
|
}
|
|
}
|
|
|
|
/// Create a minimal generation request for testing.
|
|
#[must_use]
|
|
pub fn minimal_test_request() -> GenerationRequest {
|
|
GenerationRequest {
|
|
prompt: "A simple scene".to_string(),
|
|
negative_prompt: None,
|
|
num_frames: 4,
|
|
width: 64,
|
|
height: 64,
|
|
fps: 8.0,
|
|
physics_constraints: vec![],
|
|
seed: Some(1),
|
|
guidance_scale: 5.0,
|
|
num_inference_steps: 5,
|
|
}
|
|
}
|
|
|
|
/// Create a small DiT model configuration for demos.
|
|
#[must_use]
|
|
pub fn small_dit_config() -> ModelConfig {
|
|
ModelConfig {
|
|
model_type: ModelType::DiT,
|
|
hidden_dim: 256,
|
|
num_layers: 4,
|
|
num_heads: 4,
|
|
patch_size: 4,
|
|
latent_channels: 4,
|
|
temporal_attention: true,
|
|
cross_attention: true,
|
|
}
|
|
}
|
|
|
|
/// Create a medium DiT model configuration.
|
|
#[must_use]
|
|
pub fn medium_dit_config() -> ModelConfig {
|
|
ModelConfig {
|
|
model_type: ModelType::DiT,
|
|
hidden_dim: 512,
|
|
num_layers: 8,
|
|
num_heads: 8,
|
|
patch_size: 2,
|
|
latent_channels: 4,
|
|
temporal_attention: true,
|
|
cross_attention: true,
|
|
}
|
|
}
|
|
|
|
/// Create a large DiT model configuration.
|
|
#[must_use]
|
|
pub fn large_dit_config() -> ModelConfig {
|
|
ModelConfig {
|
|
model_type: ModelType::DiT,
|
|
hidden_dim: 1024,
|
|
num_layers: 16,
|
|
num_heads: 16,
|
|
patch_size: 2,
|
|
latent_channels: 8,
|
|
temporal_attention: true,
|
|
cross_attention: true,
|
|
}
|
|
}
|
|
|
|
/// Create a DDIM scheduler configuration.
|
|
#[must_use]
|
|
pub fn ddim_scheduler_config() -> SchedulerConfig {
|
|
SchedulerConfig {
|
|
scheduler_type: SchedulerType::DDIM,
|
|
num_timesteps: 1000,
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
beta_schedule: "linear".to_string(),
|
|
}
|
|
}
|
|
|
|
/// Create a cosine scheduler configuration.
|
|
#[must_use]
|
|
pub fn cosine_scheduler_config() -> SchedulerConfig {
|
|
SchedulerConfig {
|
|
scheduler_type: SchedulerType::DDPM,
|
|
num_timesteps: 1000,
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
beta_schedule: "cosine".to_string(),
|
|
}
|
|
}
|
|
|
|
/// Create a fast scheduler configuration for testing.
|
|
#[must_use]
|
|
pub fn fast_scheduler_config() -> SchedulerConfig {
|
|
SchedulerConfig {
|
|
scheduler_type: SchedulerType::Euler,
|
|
num_timesteps: 100,
|
|
beta_start: 0.001,
|
|
beta_end: 0.02,
|
|
beta_schedule: "linear".to_string(),
|
|
}
|
|
}
|
|
|
|
/// Create a sample training configuration.
|
|
#[must_use]
|
|
pub fn sample_training_config() -> TrainingConfig {
|
|
TrainingConfig {
|
|
epochs: 100,
|
|
batch_size: 4,
|
|
learning_rate: 1e-4,
|
|
weight_decay: 0.01,
|
|
gradient_accumulation: 4,
|
|
mixed_precision: true,
|
|
physics_weight: 0.1,
|
|
seed: Some(42),
|
|
}
|
|
}
|
|
|
|
/// Create a fast training configuration for demos.
|
|
#[must_use]
|
|
pub fn fast_training_config() -> TrainingConfig {
|
|
TrainingConfig {
|
|
epochs: 10,
|
|
batch_size: 2,
|
|
learning_rate: 1e-4,
|
|
weight_decay: 0.01,
|
|
gradient_accumulation: 1,
|
|
mixed_precision: false,
|
|
physics_weight: 0.1,
|
|
seed: Some(42),
|
|
}
|
|
}
|
|
|
|
/// Create a sample world state for physics simulation.
|
|
#[must_use]
|
|
pub fn sample_world_state() -> WorldState {
|
|
WorldState {
|
|
objects: vec![
|
|
WorldObject {
|
|
id: "ball".to_string(),
|
|
object_type: "sphere".to_string(),
|
|
position: [0.0, 2.0, 0.0],
|
|
velocity: [1.0, 0.0, 0.0],
|
|
rotation: [0.0, 0.0, 0.0],
|
|
scale: [0.5, 0.5, 0.5],
|
|
mass: 1.0,
|
|
},
|
|
WorldObject {
|
|
id: "floor".to_string(),
|
|
object_type: "plane".to_string(),
|
|
position: [0.0, 0.0, 0.0],
|
|
velocity: [0.0, 0.0, 0.0],
|
|
rotation: [0.0, 0.0, 0.0],
|
|
scale: [10.0, 1.0, 10.0],
|
|
mass: f32::INFINITY,
|
|
},
|
|
WorldObject {
|
|
id: "wall".to_string(),
|
|
object_type: "box".to_string(),
|
|
position: [5.0, 1.0, 0.0],
|
|
velocity: [0.0, 0.0, 0.0],
|
|
rotation: [0.0, 0.0, 0.0],
|
|
scale: [0.5, 2.0, 5.0],
|
|
mass: f32::INFINITY,
|
|
},
|
|
],
|
|
camera_position: [5.0, 3.0, 5.0],
|
|
camera_rotation: [-0.5, 0.8, 0.0],
|
|
lighting: LightingState {
|
|
ambient: [0.1, 0.1, 0.12],
|
|
sun_direction: [0.5, -0.8, 0.3],
|
|
sun_color: [1.0, 0.95, 0.9],
|
|
time_of_day: 14.0,
|
|
},
|
|
time: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Create an underwater world state.
|
|
#[must_use]
|
|
pub fn underwater_world_state() -> WorldState {
|
|
WorldState {
|
|
objects: vec![
|
|
WorldObject {
|
|
id: "jellyfish".to_string(),
|
|
object_type: "soft_body".to_string(),
|
|
position: [0.0, 5.0, 0.0],
|
|
velocity: [0.0, -0.2, 0.0],
|
|
rotation: [0.0, 0.0, 0.0],
|
|
scale: [1.0, 1.5, 1.0],
|
|
mass: 0.5,
|
|
},
|
|
WorldObject {
|
|
id: "seabed".to_string(),
|
|
object_type: "terrain".to_string(),
|
|
position: [0.0, 0.0, 0.0],
|
|
velocity: [0.0, 0.0, 0.0],
|
|
rotation: [0.0, 0.0, 0.0],
|
|
scale: [20.0, 1.0, 20.0],
|
|
mass: f32::INFINITY,
|
|
},
|
|
],
|
|
camera_position: [3.0, 5.0, 3.0],
|
|
camera_rotation: [-0.3, 0.5, 0.0],
|
|
lighting: LightingState {
|
|
ambient: [0.05, 0.1, 0.15],
|
|
sun_direction: [0.0, -1.0, 0.0],
|
|
sun_color: [0.3, 0.5, 0.7],
|
|
time_of_day: 12.0,
|
|
},
|
|
time: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Get all sample generation requests.
|
|
#[must_use]
|
|
pub fn all_sample_requests() -> Vec<GenerationRequest> {
|
|
vec![
|
|
bouncing_ball_request(),
|
|
flowing_water_request(),
|
|
pendulum_request(),
|
|
driving_car_request(),
|
|
jellyfish_request(),
|
|
falling_leaves_request(),
|
|
]
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_bouncing_ball_request() {
|
|
let req = bouncing_ball_request();
|
|
assert!(!req.prompt.is_empty());
|
|
assert!(req.num_frames > 0);
|
|
assert!(!req.physics_constraints.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_flowing_water_request() {
|
|
let req = flowing_water_request();
|
|
assert!(req.prompt.contains("water"));
|
|
assert!(
|
|
req.physics_constraints
|
|
.iter()
|
|
.any(|c| c.constraint_type == PhysicsConstraintType::Fluid)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pendulum_request() {
|
|
let req = pendulum_request();
|
|
assert!(req.prompt.contains("pendulum"));
|
|
assert!(
|
|
req.physics_constraints
|
|
.iter()
|
|
.any(|c| c.constraint_type == PhysicsConstraintType::RigidBody)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_minimal_test_request() {
|
|
let req = minimal_test_request();
|
|
assert_eq!(req.num_frames, 4);
|
|
assert_eq!(req.width, 64);
|
|
assert_eq!(req.height, 64);
|
|
assert!(req.physics_constraints.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_configs() {
|
|
let small = small_dit_config();
|
|
let medium = medium_dit_config();
|
|
let large = large_dit_config();
|
|
|
|
assert!(small.hidden_dim < medium.hidden_dim);
|
|
assert!(medium.hidden_dim < large.hidden_dim);
|
|
assert!(small.num_layers < medium.num_layers);
|
|
assert!(medium.num_layers < large.num_layers);
|
|
}
|
|
|
|
#[test]
|
|
fn test_scheduler_configs() {
|
|
let ddim = ddim_scheduler_config();
|
|
let cosine = cosine_scheduler_config();
|
|
let fast = fast_scheduler_config();
|
|
|
|
assert_eq!(ddim.scheduler_type, SchedulerType::DDIM);
|
|
assert_eq!(cosine.beta_schedule, "cosine");
|
|
assert!(fast.num_timesteps < ddim.num_timesteps);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_configs() {
|
|
let sample = sample_training_config();
|
|
let fast = fast_training_config();
|
|
|
|
assert!(fast.epochs < sample.epochs);
|
|
assert!(fast.batch_size <= sample.batch_size);
|
|
}
|
|
|
|
#[test]
|
|
fn test_world_states() {
|
|
let world = sample_world_state();
|
|
assert!(!world.objects.is_empty());
|
|
assert!(world.objects.iter().any(|o| o.id == "ball"));
|
|
assert!(world.objects.iter().any(|o| o.id == "floor"));
|
|
|
|
let underwater = underwater_world_state();
|
|
assert!(underwater.objects.iter().any(|o| o.id == "jellyfish"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_sample_requests() {
|
|
let requests = all_sample_requests();
|
|
assert!(!requests.is_empty());
|
|
|
|
for req in requests {
|
|
assert!(!req.prompt.is_empty());
|
|
assert!(req.num_frames > 0);
|
|
assert!(req.width > 0);
|
|
assert!(req.height > 0);
|
|
}
|
|
}
|
|
}
|