From f31cf35da18a07291af6addd2b1a54fca57afea3 Mon Sep 17 00:00:00 2001 From: Danny McGee Date: Wed, 28 Feb 2024 15:52:01 -0500 Subject: [PATCH] More Image Generator stuff (#59) YOLO --- .vscode/settings.json | 5 +- image-generator-demo/project.json | 4 +- image-generator-demo/src/app/app.element.scss | 36 ++- image-generator-demo/src/app/app.element.ts | 215 +++++++++++++----- image-generator-demo/src/app/lcm-client.ts | 172 ++++++++++++++ studio-ui/src/lib/icon.element.ts | 8 + .../src/lib/scene-hierarchy.element.scss | 4 + studio/Cargo.toml | 1 - studio/project.json | 8 - studio/src/change_tracking.rs | 15 ++ studio/src/input.rs | 16 +- studio/src/interaction/gizmos/common.rs | 2 +- studio/src/interaction/gizmos/mod.rs | 1 + studio/src/lib.rs | 2 +- studio/src/scene.rs | 20 +- studio/src/wasm/mod.rs | 2 - studio/src/wasm/screenshot.rs | 68 +++++- studio/src/world_grid.rs | 2 +- 18 files changed, 462 insertions(+), 119 deletions(-) create mode 100644 image-generator-demo/src/app/lcm-client.ts diff --git a/.vscode/settings.json b/.vscode/settings.json index 4f9b6e6..a74e4ec 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -4,10 +4,7 @@ "*.element.ts": "${capture}.element.scss, ${capture}.element.spec.ts", "*.tsx": "${capture}.module.scss, ${capture}.spec.tsx", }, - "rust-analyzer.cargo.features": [ - "wasm", - "wasm_screenshot" - ], + "rust-analyzer.cargo.features": ["wasm"], "files.associations": { "*.meta": "ron" } diff --git a/image-generator-demo/project.json b/image-generator-demo/project.json index 719f5b8..396ee34 100644 --- a/image-generator-demo/project.json +++ b/image-generator-demo/project.json @@ -29,7 +29,7 @@ "dependsOn": [ { "projects": ["studio"], - "target": "build:image-generator" + "target": "build:web" } ], "options": { @@ -52,7 +52,7 @@ "dependsOn": [ { "projects": ["studio"], - "target": "build:image-generator" + "target": "build:web" } ], "options": { diff --git a/image-generator-demo/src/app/app.element.scss b/image-generator-demo/src/app/app.element.scss index 749c3dd..2687776 100644 --- a/image-generator-demo/src/app/app.element.scss +++ b/image-generator-demo/src/app/app.element.scss @@ -27,10 +27,7 @@ display: block; width: 100%; height: 100%; - padding: 1rem { - inline-end: 0.5rem; - } - border: none; + border-inline-end: none; background: transparent; } .hierarchy-resize-handle { @@ -73,8 +70,9 @@ border-radius: 0.5rem; } -.config-prompt { +.config-prompts { display: flex; + flex-direction: column; gap: 0.5rem; } @@ -118,10 +116,6 @@ color: #FFFF; } - &.prompt { - width: 512px; - } - &.seed { width: 96px; appearance: textfield; @@ -134,8 +128,32 @@ } } +.prompt { + flex: 1 0 0; +} + .config-field-group { + display: flex; + align-items: center; + gap: 0.5rem; position: relative; + + .pre-icon { + &.red { + color: rgb(220, 53, 69); + font-size: 1.333rem; + } + &.green { + color: rgb(25, 135, 84); + font-size: 1.15rem; + } + } +} + +.config-row { + display: flex; + width: 768px; + gap: 0.5rem; } .seed-refresh { diff --git a/image-generator-demo/src/app/app.element.ts b/image-generator-demo/src/app/app.element.ts index 606c239..6d18acd 100644 --- a/image-generator-demo/src/app/app.element.ts +++ b/image-generator-demo/src/app/app.element.ts @@ -7,10 +7,13 @@ import { StudioMode, } from "@storyteller/studio"; import { match } from "@storyteller/utility"; -import { LitElement, html, nothing } from "lit"; +import { LitElement, type PropertyValues, html, nothing } from "lit"; import { customElement, state } from "lit/decorators.js"; -import { styleMap } from "lit/directives/style-map.js"; import { classMap } from "lit/directives/class-map.js"; +import { createRef, ref } from "lit/directives/ref.js"; +import { styleMap } from "lit/directives/style-map.js"; + +import { ImageInfo, LcmClient } from "./lcm-client"; import "@storyteller/studio-web/canvas-provider"; import "@storyteller/studio-web/scene-hierarchy"; @@ -21,26 +24,81 @@ import "./app.element.scss" @customElement("app-root") export class AppElement extends LitElement { - @state() imageCaptureUrl?: string; - @state() realtimeCapture = true; - @state() hierarchyWidth = 256; @state() resizing = false; @state() loading = true; + @state() blob?: Blob; @state() prompt = ""; + @state() negPrompt = ""; @state() seed = randomSeed(); + @state() inputUrl?: string; + @state() outputUrl?: string; + + #lcmClient = new LcmClient(crypto.randomUUID()); + #outputImageRef = createRef(); + + constructor () { + super(); + + this.#lcmClient.addEventListener("lcm-image-ready", event => { + console.log("lcm-image-ready", event.detail); + + const prevUrl = this.outputUrl; + this.outputUrl = URL.createObjectURL(event.detail); + + if (prevUrl) + URL.revokeObjectURL(prevUrl); + }); + + this.#lcmClient.addEventListener("lcm-image-url-ready", event => { + console.log("lcm-image-url-ready", event.detail); + this.outputUrl = event.detail; + this.requestUpdate(); + }); + } + protected override createRenderRoot = () => this; override disconnectedCallback(): void { - if (this.imageCaptureUrl != null) - URL.revokeObjectURL(this.imageCaptureUrl); + if (this.inputUrl) + URL.revokeObjectURL(this.inputUrl); + + if (this.outputUrl) + URL.revokeObjectURL(this.outputUrl); super.disconnectedCallback(); } + protected override willUpdate(changes: PropertyValues): void { + let updates: Partial = {}; + + if (changes.has("blob")) + updates.blob = this.blob; + + if (changes.has("prompt")) + updates.prompt = this.prompt; + + if (changes.has("negPrompt")) + updates.negativePrompt = this.negPrompt; + + if (changes.has("seed")) + updates.seed = this.seed; + + if ( + changes.has("blob") + || changes.has("prompt") + || changes.has("negPrompt") + || changes.has("seed") + ) { + this.#lcmClient.update(updates); + } + + super.willUpdate(changes); + } + @on("scene-state") onSceneStateChange({ detail: state }: SceneStateEvent): void { console.log("scene state change:", SceneState[state]); @@ -52,12 +110,9 @@ export class AppElement extends LitElement { } @on("entity-changes") + @debounce(1000 / 24) onStudioChanges(): void { - if (this.realtimeCapture) { - this.requestScreenshot_realtime(); - } else { - this.requestScreenshot_delayed(); - } + studio.requestScreenshot(); } @on("viewport-resized") @@ -67,27 +122,40 @@ export class AppElement extends LitElement { studio.requestScreenshot(); } - @debounce(1000 / 24, { maxWait: 1000 / 5 }) - requestScreenshot_realtime() { - studio.requestScreenshot(); - } - - @debounce(1000 / 24) - requestScreenshot_delayed() { - studio.requestScreenshot(); - } - @on("screenshot-ready") onScreenshotReady({ detail: screenshot }: ScreenshotReadyEvent) { - const prevUrl = this.imageCaptureUrl; + this.blob = new Blob([screenshot.data.buffer], { type: "image/jpg" }); - const blob = new Blob([screenshot.data.buffer], { type: "image/png" }); - this.imageCaptureUrl = URL.createObjectURL(blob); + const prevUrl = this.inputUrl; + this.inputUrl = URL.createObjectURL(this.blob); if (prevUrl != null) URL.revokeObjectURL(prevUrl); } + updatePrompt(event: Event): void { + this.prompt = (event.target as HTMLInputElement).value; + } + + updateNegPrompt(event: Event): void { + this.negPrompt = (event.target as HTMLInputElement).value; + } + + refreshSeed(): void { + this.seed = randomSeed(); + } + + stopTextInputBubbling(event: KeyboardEvent): void { + // We want modifier keys like "Alt" to always be received by the Bevy + // canvas because they affect the camera controller behavior, but we do + // NOT want Bevy to receive typing inputs like W, A, S, and D while the + // user is trying to type into the Prompt field. + if (/^(Alt|Control)$/.test(event.key)) + return; + + event.stopPropagation(); + } + onResizeStart(): void { this.resizing = true; } @@ -101,15 +169,6 @@ export class AppElement extends LitElement { this.hierarchyWidth = Math.max(Math.min(updated, 512), 128); } - toggleRealtimeCapture(event: InputEvent): void { - const { checked } = event.target as HTMLInputElement; - this.realtimeCapture = checked; - } - - refreshSeed(): void { - this.seed = randomSeed(); - } - protected override render = () => html`
- ${this.imageCaptureUrl != null ? html` + ${this.inputUrl || this.outputUrl ? html` Image capture ` : html`
`} -
event.stopPropagation()} - @keyup=${(event: KeyboardEvent) => event.stopPropagation()} +
- -
+
+ { + if (event.key === "Enter") { + event.preventDefault(); + this.updatePrompt(event); + } + }} /> - +
+
+
+ + { + if (event.key === "Enter") { + event.preventDefault(); + this.updateNegPrompt(event); + } + }} + /> +
+
+ + +
+ ${/* + FIXME: Real-time capture is not really feasible until we can come up + with a more elegant way to exclude the world grid, gizmos, etc. + from the captured image +
+ */ + nothing} ${this.loading ? html` @@ -206,5 +303,5 @@ export class AppElement extends LitElement { } function randomSeed(): number { - return Math.round(Math.random() * 5000); + return Math.round(Math.random() * 99999); } diff --git a/image-generator-demo/src/app/lcm-client.ts b/image-generator-demo/src/app/lcm-client.ts new file mode 100644 index 0000000..58b156c --- /dev/null +++ b/image-generator-demo/src/app/lcm-client.ts @@ -0,0 +1,172 @@ +export enum LcmLiveStatus { + None = "NONE", + Connected = "connected", + Disconnected = "disconnected", + Wait = "wait", + SendFrame = "send_frame", + Timeout = "timeout", +} + +export interface ImageInfo { + blob: Blob; + prompt: string; + negativePrompt: string; + seed: number; + guidanceScale: number; +} + +interface LcmMessageData { + status: LcmLiveStatus | "error"; + message?: string; +} + +export class LcmClient extends WebSocket { + static readonly Api = "//b3000.pancake-chimaera.ts.net/api"; + + #userId: string; + + get status() { return this.#status; } + #status = LcmLiveStatus.None; + + #waiting = false; + + #blob?: Blob; + #prompt?: string; + #negativePrompt?: string; + #seed?: number; + #guidanceScale = 0.8; + + constructor (userId: string) { + super(`wss:${LcmClient.Api}/ws/${userId}`); + this.#userId = userId; + + this.addEventListener("open", () => this.#onOpen()); + this.addEventListener("close", () => this.#onClose()); + this.addEventListener("message", event => this.#onMessage(event)) + } + + update(data: Partial): void { + if ("blob" in data) + this.#blob = data.blob; + + if ("prompt" in data) + this.#prompt = data.prompt; + + if ("negativePrompt" in data) + this.#negativePrompt = data.negativePrompt; + + if ("seed" in data) + this.#seed = data.seed; + + this.#update(); + } + + #onOpen(): void { + console.log(`Marking connection to websocket as open for user ${this.#userId}`); + this.#status = LcmLiveStatus.Connected; + } + + #onClose(): void { + console.log(`Marking connection to websocket as closed for user ${this.#userId}`); + this.#status = LcmLiveStatus.Disconnected; + } + + #onMessage(event: MessageEvent): void { + console.log("onMessage:", event); + const data: LcmMessageData = JSON.parse(event.data); + + if (data.status === "error") { + this.#status = LcmLiveStatus.Disconnected; + console.error(data.message); + + return; + } + + this.#status = data.status; + + if (this.#status === LcmLiveStatus.SendFrame && this.#waiting) + this.#sendFrame(); + } + + #update(): void { + if (this.#status === LcmLiveStatus.SendFrame) { + this.#sendFrame(); + } else { + this.#waiting = true; + } + } + + async #sendFrame(): Promise { + if (!this.#blob || !this.#prompt || !this.#negativePrompt || this.#seed == null) { + this.#waiting = true; + return; + } + + console.log("Sending next frame..."); + + this.send(JSON.stringify({ status: "next_frame" })); + this.send(JSON.stringify({ + prompt: this.#prompt, + negativePrompt: this.#negativePrompt, + seed: this.#seed, + guidanceScale: this.#guidanceScale, + })); + this.send(this.#blob); + + this.#waiting = false; + + try { + const url = `https:${LcmClient.Api}/stream/${this.#userId}`; + + // FIXME: This only actually renders an image in the DOM about ~10% of the time + this.dispatchEvent(new CustomEvent("lcm-image-url-ready", { detail: url })); + + // const response = await fetch(url, { method: "GET" }); + // console.log("response:", response); + + // if (!response.ok || !response.body) { + // console.error(`GET failed with status ${response.status}: "${response.statusText}"`); + // return; + // } + + // const blobParts = [] as Uint8Array[]; + // const reader = response.body.getReader(); + + // // FIXME: This doesn't appear to be an image + // // const { value } = await reader.read(); + // // if (value) blobParts.push(value); + + // // FIXME: this just stops responding after the first `read()` call even + // // though `done` is false + // while (true) { + // console.log("Reading next chunk..."); + // + // const { value, done } = await reader.read(); + // console.log({ value, done }); + // + // if (value) blobParts.push(value); + // if (done) break; + // } + + // this.dispatchEvent(new CustomEvent("lcm-image-ready", { + // detail: new Blob(blobParts, { type: "image/jpg" }), + // })); + } + catch (err) { + console.error(err); + } + } +} + +export interface LcmClientEventMap extends WebSocketEventMap { + "lcm-image-ready": CustomEvent; + "lcm-image-url-ready": CustomEvent; +} + +export interface LcmClient { + addEventListener( + type: K, + listener: (this: LcmClient, ev: LcmClientEventMap[K]) => any, + options?: boolean | AddEventListenerOptions + ): void; +} diff --git a/studio-ui/src/lib/icon.element.ts b/studio-ui/src/lib/icon.element.ts index b31bcb4..0f10b2f 100644 --- a/studio-ui/src/lib/icon.element.ts +++ b/studio-ui/src/lib/icon.element.ts @@ -24,6 +24,10 @@ import fasFloppyDisk from "@fortawesome/fontawesome-pro/svgs/solid/floppy-disk.s import farSkeleton from "@fortawesome/fontawesome-pro/svgs/regular/skeleton.svg?raw"; import farBone from "@fortawesome/fontawesome-pro/svgs/regular/bone.svg?raw"; import fasBrightness from "@fortawesome/fontawesome-pro/svgs/solid/brightness.svg?raw"; +import fasCheck from "@fortawesome/fontawesome-pro/svgs/solid/check.svg?raw"; +import fasXmark from "@fortawesome/fontawesome-pro/svgs/solid/xmark.svg?raw"; +import fasCircleCheck from "@fortawesome/fontawesome-pro/svgs/solid/circle-check.svg?raw"; +import fasCircleXmark from "@fortawesome/fontawesome-pro/svgs/solid/circle-xmark.svg?raw"; import styles from "./icon.element.scss?inline"; @@ -53,7 +57,10 @@ function renderSvg(iconId: string) { "bone": () => farBone, "bwd-fast": () => fasBwdFast, "bwd-step": () => fasBwdStep, + "check": () => fasCheck, + "circle-check": () => fasCircleCheck, "circle-question": () => fasCircleQuestion, + "circle-xmark": () => fasCircleXmark, "caret-right": () => fasCaretRight, "crosshairs": () => farCrosshairs, "cube": () => farCube, @@ -70,6 +77,7 @@ function renderSvg(iconId: string) { "skeleton": () => farSkeleton, "up-down-left-right": () => farUpDownLeftRight, "video": () => farVideo, + "xmark": () => fasXmark, _: () => null, }); diff --git a/studio-web/src/lib/scene-hierarchy.element.scss b/studio-web/src/lib/scene-hierarchy.element.scss index 7cc3c00..7d504c9 100644 --- a/studio-web/src/lib/scene-hierarchy.element.scss +++ b/studio-web/src/lib/scene-hierarchy.element.scss @@ -20,5 +20,9 @@ background: #FFF3; } } + + &-corner { + background: transparent; + } } } diff --git a/studio/Cargo.toml b/studio/Cargo.toml index 8303db2..4f6b018 100644 --- a/studio/Cargo.toml +++ b/studio/Cargo.toml @@ -67,7 +67,6 @@ wasm = [ "bevy_rapier3d/wasm-bindgen", ] desktop = [] -wasm_screenshot = [] [dependencies.bevy] workspace = true diff --git a/studio/project.json b/studio/project.json index eac8000..b3b28be 100644 --- a/studio/project.json +++ b/studio/project.json @@ -31,14 +31,6 @@ "cwd": "studio" } }, - "build:image-generator": { - "executor": "nx:run-commands", - "inputs": ["assets"], - "options": { - "command": "node scripts/build-web.mjs --features=wasm,wasm_screenshot", - "cwd": "studio" - } - }, "dev": { "executor": "@nxrs/cargo:run", "options": { diff --git a/studio/src/change_tracking.rs b/studio/src/change_tracking.rs index a3f4e32..1b92f44 100644 --- a/studio/src/change_tracking.rs +++ b/studio/src/change_tracking.rs @@ -1,5 +1,7 @@ use bevy::prelude::*; +use crate::scene::{SceneElement, SceneSpawnerSystem, SceneState}; + pub struct ChangeTrackingPlugin; impl Plugin for ChangeTrackingPlugin { @@ -8,6 +10,10 @@ impl Plugin for ChangeTrackingPlugin { // Add more of these systems for each component that we want to track app.add_systems(Update, send_change_events::); + app.add_systems( + OnEnter(SceneState::Active), + track_scene_elements.after(SceneSpawnerSystem), + ); #[cfg(feature = "wasm")] app.add_systems( @@ -29,6 +35,15 @@ fn send_change_events( } } +fn track_scene_elements( + mut cmd: Commands, + q_scene_elements: Query, With, Without)>, +) { + for element in q_scene_elements.iter() { + cmd.entity(element).insert(ChangeTracker); + } +} + #[derive(Event, Deref)] pub struct EntityChangedEvent(pub Entity); diff --git a/studio/src/input.rs b/studio/src/input.rs index 32882e2..8e415de 100644 --- a/studio/src/input.rs +++ b/studio/src/input.rs @@ -177,11 +177,17 @@ fn translate_dcc_camera_events( else if (keyboard.just_released(KeyCode::AltLeft) && !keyboard.pressed(KeyCode::AltRight)) || (keyboard.just_released(KeyCode::AltRight) && !keyboard.pressed(KeyCode::AltLeft)) { - input.send(InputAction( - Lifecycle::Stop, - CameraMotion::Orbit(Vec2::ZERO).into(), - )); - input.send(InputAction(Lifecycle::Stop, CameraMotion::Zoom(0.).into())); + // Stop Camera Orbit + if mouse_buttons.pressed(MouseButton::Left) { + input.send(InputAction( + Lifecycle::Stop, + CameraMotion::Orbit(Vec2::ZERO).into(), + )); + } + // Stop Camera Zoom + else if mouse_buttons.pressed(MouseButton::Right) { + input.send(InputAction(Lifecycle::Stop, CameraMotion::Zoom(0.).into())); + } } // Camera Rotation else if mouse_buttons.pressed(MouseButton::Right) { diff --git a/studio/src/interaction/gizmos/common.rs b/studio/src/interaction/gizmos/common.rs index b507901..6c5e46f 100644 --- a/studio/src/interaction/gizmos/common.rs +++ b/studio/src/interaction/gizmos/common.rs @@ -6,7 +6,7 @@ use crate::interaction::TransformSpace; use super::material::GizmoMaterial; #[derive(Component, Reflect)] -pub(super) struct TransformGizmo { +pub struct TransformGizmo { pub(super) target: Entity, pub(super) transform_space: TransformSpace, } diff --git a/studio/src/interaction/gizmos/mod.rs b/studio/src/interaction/gizmos/mod.rs index 1a4c0fc..5ed674d 100644 --- a/studio/src/interaction/gizmos/mod.rs +++ b/studio/src/interaction/gizmos/mod.rs @@ -11,6 +11,7 @@ use self::{ rotation::{spawn_rotation_gizmo, update_ss_ring_rotation}, translation::spawn_translation_gizmo, }; +pub use common::TransformGizmo; use super::{Selection, TransformMode, TransformType}; diff --git a/studio/src/lib.rs b/studio/src/lib.rs index 0ef7ff4..8bdfb58 100644 --- a/studio/src/lib.rs +++ b/studio/src/lib.rs @@ -158,7 +158,7 @@ pub fn start(init: InitParams, mode: StudioMode, window_config: Option) ..default() }), CommonPlugins, - #[cfg(all(feature = "wasm", feature = "wasm_screenshot"))] + #[cfg(feature = "wasm")] wasm::ScreenshotPlugin, )); diff --git a/studio/src/scene.rs b/studio/src/scene.rs index 5c51325..9cd07df 100644 --- a/studio/src/scene.rs +++ b/studio/src/scene.rs @@ -23,7 +23,6 @@ use wasm_bindgen::prelude::wasm_bindgen; use wgpu::{TextureDescriptor, TextureUsages}; use crate::{ - anim::SkeletalJoint, camera_controller::{CameraController, CameraControllerBundle}, change_tracking::{ChangeTracker, ChangeTrackingPlugin}, debug::DebugPoint, @@ -63,13 +62,9 @@ impl Plugin for ScenePlugin { .add_systems( OnEnter(SceneState::Active), ( - ( - tag_scene_elements, - frame_scene_in_view, - start_capture.run_if(is_headless), - ), - apply_deferred, - track_scene_elements, + tag_scene_elements, + frame_scene_in_view, + start_capture.run_if(is_headless), ) .chain() .in_set(SceneSpawnerSystem), @@ -524,15 +519,6 @@ fn start_capture(mut captures: Query<&mut ScreenCapture>) { } } -fn track_scene_elements( - mut commands: Commands, - scene_elements: Query, With)>, -) { - for element in scene_elements.iter() { - commands.entity(element).insert(ChangeTracker); - } -} - #[cfg(feature = "wasm")] mod wasm { use bevy::{prelude::*, window::PrimaryWindow}; diff --git a/studio/src/wasm/mod.rs b/studio/src/wasm/mod.rs index c109cad..3f1770a 100644 --- a/studio/src/wasm/mod.rs +++ b/studio/src/wasm/mod.rs @@ -5,12 +5,10 @@ use wasm_bindgen::prelude::wasm_bindgen; mod conv; mod init; -#[cfg(feature = "wasm_screenshot")] mod screenshot; mod util; pub use conv::*; -#[cfg(feature = "wasm_screenshot")] pub use screenshot::ScreenshotPlugin; pub(crate) use util::*; diff --git a/studio/src/wasm/screenshot.rs b/studio/src/wasm/screenshot.rs index b4141c4..b450883 100644 --- a/studio/src/wasm/screenshot.rs +++ b/studio/src/wasm/screenshot.rs @@ -10,9 +10,11 @@ //! to the frontend by converting the `Vec` to a [`js_sys::Uint8Array`] use bevy::{prelude::*, render::view::screenshot::ScreenshotManager, window::PrimaryWindow}; -use image::{codecs::png::PngEncoder, ColorType, ImageEncoder}; +use image::{codecs::jpeg::JpegEncoder, ColorType, ImageEncoder}; use wasm_bindgen::prelude::{wasm_bindgen, JsValue}; +use crate::{interaction::gizmos::TransformGizmo, world_grid::WorldGridPlane}; + use super::{AsJsValue, StaticBuffer, WasmMessageBuffer}; static SCREENSHOT_REQUEST: WasmMessageBuffer<()> = WasmMessageBuffer::new(); @@ -22,16 +24,25 @@ pub struct ScreenshotPlugin; impl Plugin for ScreenshotPlugin { fn build(&self, app: &mut App) { + app.add_event::(); app.add_systems( Update, ( - take_screenshots.map(bevy::utils::error), + ( + prepare_screenshot, + apply_deferred, + take_screenshots.map(bevy::utils::error), + ) + .chain(), send_screenshots.map(bevy::utils::error), ), ); } } +#[derive(Event)] +pub struct ScreenshotRequestedEvent; + #[wasm_bindgen] #[derive(Clone)] pub struct Screenshot { @@ -74,21 +85,46 @@ impl AsJsValue for Screenshot { #[wasm_bindgen(js_name = requestScreenshot)] pub fn request_screenshot() { + info!("[wasm::screenshot] Screenshot request received"); SCREENSHOT_REQUEST.write_message(()); } +fn prepare_screenshot( + mut q_editor_visuals: Query<&mut Visibility, Or<(With, With)>>, + mut ew: EventWriter, +) { + if SCREENSHOT_REQUEST.take_message().is_none() { + return; + } + + info!( + "[wasm::screenshot] Hiding {} editor visuals", + q_editor_visuals.iter().count() + ); + + for mut visibility in q_editor_visuals.iter_mut() { + *visibility = Visibility::Hidden; + } + + ew.send(ScreenshotRequestedEvent); +} + fn take_screenshots( mut r_screenshot_mgr: ResMut, q_window: Query>, + mut er: EventReader, ) -> Result<(), String> { - if SCREENSHOT_REQUEST.take_message().is_none() { + if er.is_empty() { return Ok(()); } - let Ok(window_ent) = q_window.get_single() else { return Ok(()); }; + er.clear(); + + info!("[wasm::screenshot] Queueing screenshot request"); + r_screenshot_mgr .take_screenshot(window_ent, |img| { let img = match img.try_into_dynamic().map(|img| img.to_rgb8()) { @@ -112,8 +148,8 @@ fn take_screenshots( }; let src_buffer = &img.as_raw()[..buffer_len]; - let mut png_bytes = vec![]; - if let Err(err) = PngEncoder::new(&mut png_bytes).write_image( + let mut jpg_bytes = vec![]; + if let Err(err) = JpegEncoder::new_with_quality(&mut jpg_bytes, 70).write_image( src_buffer, img.width(), img.height(), @@ -124,7 +160,7 @@ fn take_screenshots( } SCREENSHOT.write_message(ScreenshotNative { - data: png_bytes, + data: jpg_bytes, width: img.width() as usize, height: img.height() as usize, }); @@ -146,7 +182,10 @@ declare global { } "#; -fn send_screenshots(q_window: Query<&Window, With>) -> Result<(), String> { +fn send_screenshots( + q_window: Query<&Window, With>, + mut q_editor_visuals: Query<&mut Visibility, Or<(With, With)>>, +) -> Result<(), String> { let Some(screenshot) = SCREENSHOT.take_message() else { return Ok(()); }; @@ -157,5 +196,16 @@ fn send_screenshots(q_window: Query<&Window, With>) -> Result<(), let screenshot = Screenshot::from(screenshot); - super::dispatch_to_js(window, "screenshot-ready", &screenshot) + super::dispatch_to_js(window, "screenshot-ready", &screenshot)?; + + info!( + "[wasm::screenshot] Revealing {} previously hidden editor visuals", + q_editor_visuals.iter().count() + ); + + for mut visibility in q_editor_visuals.iter_mut() { + *visibility = Visibility::Inherited; + } + + Ok(()) } diff --git a/studio/src/world_grid.rs b/studio/src/world_grid.rs index 8146db4..2c3256e 100644 --- a/studio/src/world_grid.rs +++ b/studio/src/world_grid.rs @@ -17,7 +17,7 @@ impl Plugin for WorldGridPlugin { } #[derive(Component)] -struct WorldGridPlane; +pub struct WorldGridPlane; #[derive(Asset, AsBindGroup, TypePath, Clone)] struct ProceduralGridMaterial {}