More Image Generator stuff (#59)

YOLO

.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"
 	}

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": {

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 {

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<HTMLImageElement>();
+
+	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<this>): void {
+		let updates: Partial<ImageInfo> = {};
+
+		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`
 		<sts-canvas-provider class="app-root">
 			<div
@@ -145,45 +204,81 @@ export class AppElement extends LitElement {
 					></sts-transform-toolbar>
 				</sts-studio>
 
-				${this.imageCaptureUrl != null ? html`
+				${this.inputUrl || this.outputUrl ? html`
 					<img
+						${ref(this.#outputImageRef)}
 						class="image-capture"
-						src=${this.imageCaptureUrl}
+						src=${(this.outputUrl ?? this.inputUrl)!}
 						alt="Image capture"
 					/>
 				` : html`
 					<div class="image-capture placeholder"></div>
 				`}
 
-				<div class="config config-prompt h-center v-end"
-					@keydown=${(event: KeyboardEvent) => event.stopPropagation()}
-					@keyup=${(event: KeyboardEvent) => event.stopPropagation()}
+				<div class="config config-prompts h-center v-end"
+					@keydown=${this.stopTextInputBubbling}
+					@keyup=${this.stopTextInputBubbling}
 				>
-					<input
-						class="config-field prompt"
-						type="text"
-						aria-label="Prompt"
-						placeholder="Enter a prompt to customize the output"
-					/>
-					<div class="config-field-group">
+					<div class="config-field-group prompt">
+						<sts-icon class="pre-icon green" icon="check"></sts-icon>
 						<input
-							class="config-field seed"
-							type="number"
-							aria-label="Seed"
-							value=${this.seed}
-							min="100"
-							max="10000"
+							class="config-field prompt"
+							type="text"
+							aria-label="Affirmative Prompt"
+							placeholder="Enter an affirmative prompt to customize the output"
+							.value=${this.prompt}
+							@change=${this.updatePrompt}
+							@keydown=${(event: KeyboardEvent) => {
+								if (event.key === "Enter") {
+									event.preventDefault();
+									this.updatePrompt(event);
+								}
+							}}
 						/>
-						<button
-							class="seed-refresh"
-							@click=${this.refreshSeed}
-						>
-							<sts-icon icon="rotate"></sts-icon>
-						</button>
+					</div>
+					<div class="config-row">
+						<div class="config-field-group prompt">
+							<sts-icon class="pre-icon red" icon="xmark"></sts-icon>
+							<input
+								class="config-field prompt"
+								type="text"
+								aria-label="Negative Prompt"
+								placeholder="Enter a negative prompt to customize the output"
+								.value=${this.negPrompt}
+								@change=${this.updateNegPrompt}
+								@keydown=${(event: KeyboardEvent) => {
+									if (event.key === "Enter") {
+										event.preventDefault();
+										this.updateNegPrompt(event);
+									}
+								}}
+							/>
+						</div>
+						<div class="config-field-group">
+							<input
+								class="config-field seed"
+								type="number"
+								aria-label="Seed"
+								.value=${this.seed.toString()}
+								min="1"
+								max="99999"
+							/>
+							<button
+								class="seed-refresh"
+								@click=${this.refreshSeed}
+							>
+								<sts-icon icon="rotate"></sts-icon>
+							</button>
+						</div>
 					</div>
 				</div>
 			</div>
 
+			${/*
+			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
+
 			<div class="config v-start h-end">
 				<label class="config-checkbox">
 					<input
@@ -194,6 +289,8 @@ export class AppElement extends LitElement {
 					Realtime Capture
 				</label>
 			</div>
+			*/
+			nothing}
 
 			${this.loading ? html`
 				<!-- TODO: Animated spinner or something -- talk to our designers -->
@@ -206,5 +303,5 @@ export class AppElement extends LitElement {
 }
 
 function randomSeed(): number {
-	return Math.round(Math.random() * 5000);
+	return Math.round(Math.random() * 99999);
 }

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<ImageInfo>): 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<string>): 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<void> {
+		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<Blob>("lcm-image-ready", {
+			// 	detail: new Blob(blobParts, { type: "image/jpg" }),
+			// }));
+		}
+		catch (err) {
+			console.error(err);
+		}
+	}
+}
+
+export interface LcmClientEventMap extends WebSocketEventMap {
+	"lcm-image-ready": CustomEvent<Blob>;
+	"lcm-image-url-ready": CustomEvent<string>;
+}
+
+export interface LcmClient {
+	addEventListener<K extends keyof LcmClientEventMap>(
+		type: K,
+		listener: (this: LcmClient, ev: LcmClientEventMap[K]) => any,
+		options?: boolean | AddEventListenerOptions
+	): void;
+}

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,
 	});
 

studio-web/src/lib/scene-hierarchy.element.scss

@@ -20,5 +20,9 @@
 				background: #FFF3;
 			}
 		}
+
+		&-corner {
+			background: transparent;
+		}
 	}
 }

studio/Cargo.toml

@@ -67,7 +67,6 @@ wasm = [
 	"bevy_rapier3d/wasm-bindgen",
 ]
 desktop = []
-wasm_screenshot = []
 
 [dependencies.bevy]
 workspace = true

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": {

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::<Transform>);
+        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<T: Component>(
     }
 }
 
+fn track_scene_elements(
+    mut cmd: Commands,
+    q_scene_elements: Query<Entity, (With<Transform>, With<SceneElement>, Without<ChangeTracker>)>,
+) {
+    for element in q_scene_elements.iter() {
+        cmd.entity(element).insert(ChangeTracker);
+    }
+}
+
 #[derive(Event, Deref)]
 pub struct EntityChangedEvent(pub Entity);
 

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) {

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,
 }

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};
 

studio/src/lib.rs

@@ -158,7 +158,7 @@ pub fn start(init: InitParams, mode: StudioMode, window_config: Option<Window>)
             ..default()
         }),
         CommonPlugins,
-        #[cfg(all(feature = "wasm", feature = "wasm_screenshot"))]
+        #[cfg(feature = "wasm")]
         wasm::ScreenshotPlugin,
     ));
 

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<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

@@ -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::*;
 

studio/src/wasm/screenshot.rs

@@ -10,9 +10,11 @@
 //!    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 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::<ScreenshotRequestedEvent>();
         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<WorldGridPlane>, With<TransformGizmo>)>>,
+    mut ew: EventWriter<ScreenshotRequestedEvent>,
+) {
+    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<ScreenshotManager>,
     q_window: Query<Entity, With<PrimaryWindow>>,
+    mut er: EventReader<ScreenshotRequestedEvent>,
 ) -> 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<PrimaryWindow>>) -> Result<(), String> {
+fn send_screenshots(
+    q_window: Query<&Window, With<PrimaryWindow>>,
+    mut q_editor_visuals: Query<&mut Visibility, Or<(With<WorldGridPlane>, With<TransformGizmo>)>>,
+) -> Result<(), String> {
     let Some(screenshot) = SCREENSHOT.take_message() else {
         return Ok(());
     };
@@ -157,5 +196,16 @@ fn send_screenshots(q_window: Query<&Window, With<PrimaryWindow>>) -> 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(())
 }

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 {}