Bootstrap Image Generator application (#53)

This PR begins scaffolding the Stable Diffusion Turbo application (title
TBD), which is a separate frontend for Storyteller Engine which focuses
on image generation.

### Design Overview

* The application renders a "split-screen" view, with Storyteller Studio
  on the left, and the Stable Diffusion Turbo output on the right
* When changes are made in the Studio editor:
  - A screenshot is taken from the Bevy application
  - Screenshot is piped into Stable Diffusion Turbo model
  - Model output is rendered to the screen
* This all happens in as close to real-time as possible

### Work completed in this PR

* The `studio-frontend` project has been (somewhat crudely) refactored
  to split most of its components into reusable libraries
  - `studio-ui` is a use-case-agnostic design system library for core UI
    components. These are "dumb" components which are controlled
    via `@property` inputs and send `CustomEvent`s in response to user
    input.
  - `studio-frontend` is a collection of reusable Storyteller Studio
    widgets. These are "smart" components which depend on
    `@storyteller/studio` directly, manage their own state, and require very
    little configuration or styling. They are designed to be drop-in
    building blocks that can be reused between various frontends.
* The new application has been bootstrapped as `image-generator-demo`
  - Currently, this simply grabs screenshots from the Bevy application and
    renders those screenshots on the right side of the split-screen view. A
    future PR will wire up to the Storyteller.ai backend to pipe these
    screenshots through Stable Diffusion Turbo, and render that model's
    output instead.

### Testing

* Install dependencies as described in the repo's README
* Run `npx nx serve image-generator-demo`
* Open [localhost:4200](http://localhost:4200) in your browser
  - Use the camera controller in the left side of the view to manipulate
    the Bevy scene
  - Verify the right side of the view updates with a fresh screenshot at a
    rate of ~5 fps
* Toggle the "Realtime Capture" checkbox
  - Verify the right side view no longer updates while the user is
    actively manipulating the editor
  - This should allow the app to run smoother on lower-spec hardware

---------

Co-authored-by: Danny McGee <dannymcgee@gmail.com>

studio/src/camera_controller.rs

@@ -10,7 +10,7 @@ use crate::{
     interaction::Selected,
     math::{
         Easing, InterpPosition, InterpVelocity, Interpolation, InterpolationSystem, NearlyEq,
-        Velocity,
+        NearlyZero, Velocity,
     },
 };
 
@@ -78,9 +78,6 @@ fn camera_control_kinematic(
         return;
     }
 
-    let (mut controller, mut xform, proj) = q_camera.single_mut();
-    let mut window = q_window.single_mut();
-
     let mut any_ongoing = false;
     let mut any_started = false;
     let mut any_stopped = false;
@@ -92,7 +89,8 @@ fn camera_control_kinematic(
         )
     }) {
         if let InputAction(lifecycle, Action::Camera(motion)) = action {
-            use CameraMotion::*;
+            let (mut controller, mut xform, proj) = q_camera.single_mut();
+            let window = q_window.single();
 
             match lifecycle {
                 Lifecycle::Start => any_started = true,
@@ -100,20 +98,21 @@ fn camera_control_kinematic(
                 Lifecycle::Stop => any_stopped = true,
             }
 
+            use CameraMotion::*;
             match motion {
-                Pan(delta) => do_pan(proj, &window, *delta, &mut xform, &mut controller),
-                Zoom(delta) => do_zoom(&window, *delta, &mut xform, &mut controller),
-                Orbit(delta) => do_orbit(&window, *delta, *lifecycle, &mut xform, &mut controller),
-                Rotate(delta) => do_rotate(&window, *delta, &mut xform, &mut controller),
+                Pan(delta) => do_pan(proj, window, *delta, &mut xform, &mut controller),
+                Zoom(delta) => do_zoom(window, *delta, &mut xform, &mut controller),
+                Orbit(delta) => do_orbit(window, *delta, *lifecycle, &mut xform, &mut controller),
+                Rotate(delta) => do_rotate(window, *delta, &mut xform, &mut controller),
                 _ => {}
             }
         }
     }
 
     if any_started {
-        capture_pointer(&mut window);
+        capture_pointer(&mut q_window.single_mut());
     } else if any_stopped && !any_ongoing {
-        release_pointer(&mut window);
+        release_pointer(&mut q_window.single_mut());
     }
 }
 
@@ -215,12 +214,16 @@ fn process_focus_inputs(
 }
 
 fn camera_control_physical(
-    mut q_camera: Query<(&mut CameraController, &mut Transform, &Velocity)>,
+    mut q_camera: Query<(Entity, &mut CameraController, &Velocity)>,
+    mut q_xforms: Query<&mut Transform>,
 ) {
-    for (mut controller, mut xform, velocity) in q_camera.iter_mut() {
-        let delta_pos = xform.rotation * (**velocity);
-        xform.translation += delta_pos;
-        controller.focus += delta_pos;
+    for (ent, mut controller, velocity) in q_camera.iter_mut() {
+        if !velocity.nearly_zero() {
+            let mut xform = q_xforms.get_mut(ent).unwrap();
+            let delta_pos = xform.rotation * (**velocity);
+            xform.translation += delta_pos;
+            controller.focus += delta_pos;
+        }
     }
 }
 

studio/src/change_tracking.rs

@@ -0,0 +1,77 @@
+use bevy::prelude::*;
+
+pub struct ChangeTrackingPlugin;
+
+impl Plugin for ChangeTrackingPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_event::<EntityChangedEvent>();
+
+        // Add more of these systems for each component that we want to track
+        app.add_systems(Update, send_change_events::<Transform>);
+
+        #[cfg(feature = "wasm")]
+        app.add_systems(
+            Update,
+            wasm::send_change_events::<Transform>.map(bevy::utils::error),
+        );
+    }
+}
+
+#[derive(Component, Debug)]
+pub struct ChangeTracker;
+
+fn send_change_events<T: Component>(
+    query: Query<Entity, (Changed<T>, With<ChangeTracker>)>,
+    mut entity_change_writer: EventWriter<EntityChangedEvent>,
+) {
+    for entity in query.iter() {
+        entity_change_writer.send(EntityChangedEvent(entity));
+    }
+}
+
+#[derive(Event, Deref)]
+pub struct EntityChangedEvent(pub Entity);
+
+#[cfg(feature = "wasm")]
+mod wasm {
+    use bevy::{prelude::*, window::PrimaryWindow};
+    use wasm_bindgen::prelude::wasm_bindgen;
+
+    use crate::wasm;
+
+    use super::ChangeTracker;
+
+    #[wasm_bindgen(typescript_custom_section)]
+    const ENTITY_CHANGES_EVENT: &str = r#"
+export interface EntityChangesEvent extends CustomEvent {
+    type: "entity-changes";
+    detail: string[];
+}
+
+declare global {
+    export interface GlobalEventHandlersEventMap {
+        "entity-changes": EntityChangesEvent;
+    }
+}
+    "#;
+
+    pub(super) fn send_change_events<T: Component>(
+        q_changed_ents: Query<Entity, (Changed<T>, With<ChangeTracker>)>,
+        q_window: Query<&Window, With<PrimaryWindow>>,
+    ) -> Result<(), String> {
+        if q_changed_ents.is_empty() {
+            return Ok(());
+        }
+
+        let Ok(window) = q_window.get_single() else {
+            return Ok(());
+        };
+
+        let entities = q_changed_ents
+            .iter()
+            .map(|ent| format!("{ent:?}"))
+            .collect::<Vec<_>>();
+
+        wasm::dispatch_to_js(window, "entity-changes", &entities)
+    }
+}

studio/src/lib.rs

@@ -19,6 +19,7 @@ use wgpu::Extent3d;
 mod anim;
 mod bvh;
 mod camera_controller;
+mod change_tracking;
 mod debug;
 mod design_system;
 mod hierarchy;
@@ -279,6 +280,8 @@ pub fn start_internal(
             ..default()
         }),
         CommonPlugins,
+        #[cfg(all(feature = "wasm", feature = "wasm_screenshot"))]
+        wasm::ScreenshotPlugin,
     ));
 
     if matches!(studio_mode, StudioMode::Editor) {

studio/src/scene.rs

@@ -22,7 +22,9 @@ use wasm_bindgen::prelude::wasm_bindgen;
 use wgpu::{TextureDescriptor, TextureUsages};
 
 use crate::{
+    anim::SkeletalJoint,
     camera_controller::{CameraController, CameraControllerBundle},
+    change_tracking::{ChangeTracker, ChangeTrackingPlugin},
     debug::DebugPoint,
     is_headless, is_not_headless,
     math::{Combinable, Extrema},
@@ -37,6 +39,8 @@ pub struct SceneSpawnerSystem;
 
 impl Plugin for ScenePlugin {
     fn build(&self, app: &mut App) {
+        app.add_plugins(ChangeTrackingPlugin);
+
         app.insert_resource(ClearColor(Color::hsl(0.0, 0.0, 0.25)))
             .insert_resource(Msaa::Off)
             .add_state::<SceneState>()
@@ -56,10 +60,15 @@ impl Plugin for ScenePlugin {
             .add_systems(
                 OnEnter(SceneState::Active),
                 (
-                    tag_scene_elements,
-                    frame_scene_in_view,
-                    start_capture.run_if(is_headless),
+                    (
+                        tag_scene_elements,
+                        frame_scene_in_view,
+                        start_capture.run_if(is_headless),
+                    ),
+                    apply_deferred,
+                    track_scene_elements,
                 )
+                    .chain()
                     .in_set(SceneSpawnerSystem),
             );
 
@@ -234,6 +243,7 @@ fn initialize(mut cmd: Commands) {
     cmd.spawn((
         Name::new("Main Camera"),
         MainCamera,
+        ChangeTracker,
         Camera3dBundle {
             transform: camera_xform,
             tonemapping: Tonemapping::BlenderFilmic,
@@ -483,6 +493,15 @@ fn start_capture(mut captures: Query<&mut ScreenCapture>) {
     }
 }
 
+fn track_scene_elements(
+    mut commands: Commands,
+    scene_elements: Query<Entity, (With<Transform>, With<SkeletalJoint>)>,
+) {
+    for element in scene_elements.iter() {
+        commands.entity(element).insert(ChangeTracker);
+    }
+}
+
 #[cfg(feature = "wasm")]
 mod wasm {
     use bevy::{prelude::*, window::PrimaryWindow};

studio/src/wasm/mod.rs

@@ -4,9 +4,13 @@ use parking_lot::ReentrantMutex;
 use wasm_bindgen::prelude::wasm_bindgen;
 
 mod conv;
+#[cfg(feature = "wasm_screenshot")]
+mod screenshot;
 mod util;
 
 pub use conv::*;
+#[cfg(feature = "wasm_screenshot")]
+pub use screenshot::ScreenshotPlugin;
 pub(crate) use util::*;
 
 static INITIALIZED: OnceLock<ReentrantMutex<RefCell<bool>>> = OnceLock::new();

studio/src/wasm/screenshot.rs

@@ -0,0 +1,161 @@
+//! Outline of the flow of events here:
+//!
+//! 1. [`crate::change_tracking`] system sends events to frontend to notify of
+//!    tracked changes
+//! 2. Frontend calls [`request_screenshot`] on a debounce, writing an empty
+//!    message to [`SCREENSHOT_REQUEST`]
+//! 3. [`take_screenshots`] watches the [`SCREENSHOT_REQUEST`] buffer and
+//!    populates the [`SCREENSHOT`] buffer with a PNG-encoded `Vec<u8>`
+//! 4. [`send_screenshots`] watches the [`SCREENSHOT`] buffer and sends an event
+//!    to the frontend by converting the `Vec<u8>` to a [`js_sys::Uint8Array`]
+
+use bevy::{prelude::*, render::view::screenshot::ScreenshotManager, window::PrimaryWindow};
+use image::{codecs::png::PngEncoder, ColorType, ImageEncoder};
+use wasm_bindgen::prelude::{wasm_bindgen, JsValue};
+
+use super::{AsJsValue, StaticBuffer, WasmMessageBuffer};
+
+static SCREENSHOT_REQUEST: WasmMessageBuffer<()> = WasmMessageBuffer::new();
+static SCREENSHOT: WasmMessageBuffer<ScreenshotNative> = WasmMessageBuffer::new();
+
+pub struct ScreenshotPlugin;
+
+impl Plugin for ScreenshotPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_systems(
+            Update,
+            (
+                take_screenshots.map(bevy::utils::error),
+                send_screenshots.map(bevy::utils::error),
+            ),
+        );
+    }
+}
+
+#[wasm_bindgen]
+#[derive(Clone)]
+pub struct Screenshot {
+    data_pvt: js_sys::Uint8Array,
+    pub width: usize,
+    pub height: usize,
+}
+
+// This intermediate type is necessary because the `Screenshot` struct
+// cannot be safely stored in a `WasmMessageBuffer`
+struct ScreenshotNative {
+    data: Vec<u8>,
+    width: usize,
+    height: usize,
+}
+
+impl From<ScreenshotNative> for Screenshot {
+    fn from(value: ScreenshotNative) -> Self {
+        Self {
+            data_pvt: js_sys::Uint8Array::from(&value.data[..]),
+            width: value.width,
+            height: value.height,
+        }
+    }
+}
+
+#[wasm_bindgen]
+impl Screenshot {
+    #[wasm_bindgen(getter)]
+    pub fn data(&self) -> js_sys::Uint8Array {
+        self.data_pvt.clone()
+    }
+}
+
+impl AsJsValue for Screenshot {
+    fn as_js(&self) -> Result<JsValue, String> {
+        Ok(JsValue::from(self.clone()))
+    }
+}
+
+#[wasm_bindgen(js_name = requestScreenshot)]
+pub fn request_screenshot() {
+    SCREENSHOT_REQUEST.write_message(());
+}
+
+fn take_screenshots(
+    mut r_screenshot_mgr: ResMut<ScreenshotManager>,
+    q_window: Query<Entity, With<PrimaryWindow>>,
+) -> Result<(), String> {
+    if SCREENSHOT_REQUEST.take_message().is_none() {
+        return Ok(());
+    }
+
+    let Ok(window_ent) = q_window.get_single() else {
+        return Ok(());
+    };
+
+    r_screenshot_mgr
+        .take_screenshot(window_ent, |img| {
+            let img = match img.try_into_dynamic().map(|img| img.to_rgb8()) {
+                Ok(img) => img,
+                Err(err) => {
+                    error!("{err}");
+                    return;
+                }
+            };
+
+            let channel_count = <image::Rgb<u8> as image::Pixel>::CHANNEL_COUNT as usize;
+            let buffer_len = match Some(channel_count)
+                .and_then(|size| size.checked_mul(img.width() as usize))
+                .and_then(|size| size.checked_mul(img.height() as usize))
+            {
+                Some(size) => size,
+                None => {
+                    error!("Failed to compute buffer length for image!");
+                    return;
+                }
+            };
+            let src_buffer = &img.as_raw()[..buffer_len];
+
+            let mut png_bytes = vec![];
+            if let Err(err) = PngEncoder::new(&mut png_bytes).write_image(
+                src_buffer,
+                img.width(),
+                img.height(),
+                ColorType::Rgb8,
+            ) {
+                error!("{err}");
+                return;
+            }
+
+            SCREENSHOT.write_message(ScreenshotNative {
+                data: png_bytes,
+                width: img.width() as usize,
+                height: img.height() as usize,
+            });
+        })
+        .map_err(|err| format!("{err}"))
+}
+
+#[wasm_bindgen(typescript_custom_section)]
+const SCREENSHOT_READY_EVENT: &str = r#"
+export interface ScreenshotReadyEvent extends CustomEvent {
+    type: "screenshot-ready";
+    detail: Screenshot;
+}
+
+declare global {
+    export interface GlobalEventHandlersEventMap {
+        "screenshot-ready": ScreenshotReadyEvent;
+    }
+}
+    "#;
+
+fn send_screenshots(q_window: Query<&Window, With<PrimaryWindow>>) -> Result<(), String> {
+    let Some(screenshot) = SCREENSHOT.take_message() else {
+        return Ok(());
+    };
+
+    let Ok(window) = q_window.get_single() else {
+        return Ok(());
+    };
+
+    let screenshot = Screenshot::from(screenshot);
+
+    super::dispatch_to_js(window, "screenshot-ready", &screenshot)
+}