Spaces:
Running on Zero
Running on Zero
Bobby
Claude Opus 5.5
Turn Generate 3D Model into a Download 3D Model button once the GLB is ready
0ce4985 | import inspect | |
| import os | |
| import pickle | |
| import sys | |
| import tempfile | |
| import time | |
| import unittest | |
| from pathlib import Path | |
| from unittest import mock | |
| ROOT = Path(__file__).resolve().parent.parent | |
| sys.path.insert(0, str(ROOT)) | |
| sys.path.insert(0, str(ROOT / "tests")) | |
| os.environ["LABS_FALLBACK_SIGNING_KEY"] = "unit-signing-key" | |
| os.environ.pop("SPACE_ID", None) | |
| os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False") | |
| import fallback_stubs # noqa: E402 | |
| fallback_stubs.install() | |
| import gradio as gr # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import trimesh # noqa: E402 | |
| from PIL import Image # noqa: E402 | |
| import app # noqa: E402 | |
| PAYLOAD = {"glb_name": "model.glb", "glb_b64": "aGFybmVzcw=="} | |
| SETTINGS = {"simplify": app.GLB_EXPORT_SIMPLIFY, "texture_size": app.GLB_EXPORT_TEXTURE_SIZE} | |
| def cached_state(**overrides): | |
| state = { | |
| "runtime": "original", | |
| "marker": "state-0", | |
| app.GLB_MESH_KEY: {"glb": "prepared"}, | |
| app.GLB_SETTINGS_KEY: dict(SETTINGS), | |
| app.GLB_PAYLOAD_KEY: dict(PAYLOAD), | |
| } | |
| state.update(overrides) | |
| return state | |
| class GlbBake(unittest.TestCase): | |
| def setUp(self): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| fallback_stubs.GPU_GATE["error"] = None | |
| def test_bake_prepares_the_mesh_at_the_requested_settings(self): | |
| prepared = object() | |
| with mock.patch.object(app, "_prepare_glb", return_value=prepared) as prepare: | |
| glb, failure = app.bake_glb_on_gpu({"runtime": "original"}, app.GLB_EXPORT_SIMPLIFY, | |
| app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| prepare.assert_called_once_with({"runtime": "original"}, app.GLB_EXPORT_TEXTURE_SIZE) | |
| self.assertIs(glb, prepared) | |
| self.assertIsNone(failure) | |
| def test_original_bake_reads_the_cpu_packed_state(self): | |
| from types import SimpleNamespace | |
| from export_profiles import simplify_ratio_for_texture_size | |
| source_faces = 400000 | |
| mesh = SimpleNamespace(faces=[0] * source_faces) | |
| gaussian, prepared = object(), object() | |
| state = {"runtime": "original", "mesh": {"packed": True}, "gaussian": {"packed": True}} | |
| with mock.patch.object(app, "HAS_DIFF_GAUSSIAN_RASTERIZATION", True), \ | |
| mock.patch.object(app, "unpack_original_state", | |
| return_value={"mesh": mesh, "gaussian": gaussian}), \ | |
| mock.patch.object(app.trellis_postprocessing_utils, "to_glb", | |
| return_value=prepared) as convert: | |
| glb, failure = app.bake_glb_on_gpu(state, app.GLB_EXPORT_SIMPLIFY, | |
| app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| self.assertIs(convert.call_args.args[0], gaussian) | |
| self.assertIs(convert.call_args.args[1], mesh) | |
| self.assertIs(glb, prepared) | |
| self.assertIsNone(failure) | |
| self.assertEqual(convert.call_args.kwargs["texture_size"], app.GLB_EXPORT_TEXTURE_SIZE) | |
| self.assertEqual( | |
| convert.call_args.kwargs["simplify"], | |
| simplify_ratio_for_texture_size(app.GLB_EXPORT_TEXTURE_SIZE, source_faces), | |
| ) | |
| def test_bake_failure_is_reported_without_raising(self): | |
| with mock.patch.object(app, "_prepare_glb", side_effect=RuntimeError("out of memory")): | |
| glb, failure = app.bake_glb_on_gpu({"runtime": "original"}, app.GLB_EXPORT_SIMPLIFY, | |
| app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| self.assertIsNone(glb) | |
| self.assertEqual(failure, app.GLB_PREP_FAILURE_MESSAGE) | |
| def test_bake_keeps_the_state_so_the_user_can_retry(self): | |
| with mock.patch.object(app, "_prepare_glb", side_effect=RuntimeError("out of memory")): | |
| state = {"runtime": "original", "marker": "state-0"} | |
| app.bake_glb_on_gpu(state, app.GLB_EXPORT_SIMPLIFY, app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| self.assertEqual(state, {"runtime": "original", "marker": "state-0"}) | |
| self.assertNotIn(app.GLB_MESH_KEY, state) | |
| class GlbCaching(unittest.TestCase): | |
| def setUp(self): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| fallback_stubs.GPU_GATE["error"] = None | |
| def test_packaging_stores_payload_and_drops_prepared_mesh(self): | |
| state = cached_state(**{app.GLB_MESH_KEY: {"glb": "prepared"}, app.GLB_PAYLOAD_KEY: None}) | |
| with mock.patch.object(app, "save_export_payload", return_value=dict(PAYLOAD)) as save: | |
| result = app.cache_prepared_glb(state) | |
| self.assertIs(result, state) | |
| self.assertEqual(state[app.GLB_PAYLOAD_KEY], PAYLOAD) | |
| self.assertNotIn(app.GLB_MESH_KEY, state) | |
| self.assertNotIn(app.GLB_ERROR_KEY, state) | |
| self.assertIsInstance(save.call_args.args[0], dict) | |
| def test_packaging_failure_is_reported_as_a_status_value(self): | |
| state = cached_state(**{app.GLB_MESH_KEY: {"glb": "prepared"}, app.GLB_PAYLOAD_KEY: None}) | |
| error = gr.Error("Model was exported but failed validation for viewer rendering.") | |
| with mock.patch.object(app, "save_export_payload", side_effect=error): | |
| app.cache_prepared_glb(state) | |
| self.assertEqual(state[app.GLB_ERROR_KEY], "Model was exported but failed validation for viewer rendering.") | |
| self.assertNotIn(app.GLB_PAYLOAD_KEY, state) | |
| def test_state_without_a_prepared_mesh_is_untouched(self): | |
| state = {"runtime": "original", "marker": "state-0"} | |
| with mock.patch.object(app, "save_export_payload") as save: | |
| self.assertIs(app.cache_prepared_glb(state), state) | |
| save.assert_not_called() | |
| self.assertIsNone(app.cache_prepared_glb(None)) | |
| def test_payload_requires_matching_export_settings(self): | |
| state = cached_state() | |
| size = app.GLB_EXPORT_TEXTURE_SIZE | |
| self.assertEqual(app.cached_glb_payload(state, 0.95, size), PAYLOAD) | |
| self.assertIsNone(app.cached_glb_payload(state, 0.98, size)) | |
| self.assertIsNone(app.cached_glb_payload(state, 0.95, 2048)) | |
| self.assertIsNone(app.cached_glb_payload(state, 0.95, 1536)) | |
| self.assertIsNone(app.cached_glb_payload({app.GLB_PAYLOAD_KEY: dict(PAYLOAD)}, 0.95, size)) | |
| self.assertIsNone(app.cached_glb_payload(None, 0.95, size)) | |
| self.assertIsNone(app.cached_glb_payload({}, 0.95, size)) | |
| class GlbButtonSwap(unittest.TestCase): | |
| def test_success_swaps_generate_for_a_download_of_the_glb(self): | |
| generate, status, download = app.finish_glb_export("/tmp/model.glb", {"marker": 1}, "") | |
| self.assertFalse(generate["visible"]) | |
| self.assertFalse(status["visible"]) | |
| self.assertEqual((download["value"], download["visible"]), ("/tmp/model.glb", True)) | |
| def test_failure_keeps_generate_retryable_and_hides_the_download(self): | |
| generate, status, download = app.finish_glb_export(None, {"marker": 1}, "Export failed.") | |
| self.assertEqual((generate["interactive"], generate["visible"]), (True, True)) | |
| self.assertEqual((status["value"], status["visible"]), ("Export failed.", True)) | |
| self.assertEqual((download["value"], download["visible"]), (None, False)) | |
| def test_new_generation_restores_generate_and_clears_the_download(self): | |
| outputs = app.reset_generation_outputs() | |
| self.assertEqual(outputs[7], {"interactive": False, "visible": True, "__type__": "update"}) | |
| self.assertEqual((outputs[-1]["value"], outputs[-1]["visible"]), (None, False)) | |
| class ExtractModelBoundary(unittest.TestCase): | |
| def setUp(self): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| fallback_stubs.GPU_GATE["error"] = None | |
| def extract(self, state, mesh_simplify=0.95, texture_size=None): | |
| texture_size = app.GLB_EXPORT_TEXTURE_SIZE if texture_size is None else texture_size | |
| with mock.patch.object(app, "_prepare_glb", return_value={"glb": "baked"}) as prepare, \ | |
| mock.patch.object(app, "save_export_payload", return_value=dict(PAYLOAD)): | |
| result = app.extract_model(state, mesh_simplify, texture_size, None) | |
| return result, prepare | |
| def test_cached_payload_is_returned_without_conversion_or_allocation(self): | |
| (status, payload), prepare = self.extract(cached_state()) | |
| self.assertEqual(payload, PAYLOAD) | |
| self.assertIn("Finalizing", status["value"]) | |
| prepare.assert_not_called() | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, []) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, []) | |
| def test_absent_state_returns_an_empty_update(self): | |
| (status, payload), prepare = self.extract(None) | |
| self.assertIsNone(payload) | |
| self.assertNotIn("value", status) | |
| prepare.assert_not_called() | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, []) | |
| def test_a_later_size_choice_bakes_under_a_second_allocation(self): | |
| (status, payload), prepare = self.extract(cached_state(), texture_size=2048) | |
| self.assertEqual(payload, PAYLOAD) | |
| self.assertEqual(prepare.call_args.args[1:], (2048,)) | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, ["gpu_task:bake"]) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, ["gpu_task:bake"]) | |
| self.assertNotIn("visible", status) | |
| def test_settings_choice_is_recorded_only_once_the_bake_succeeds(self): | |
| state = cached_state(**{app.GLB_PAYLOAD_KEY: None, | |
| app.GLB_SETTINGS_KEY: {"simplify": 0.95, "texture_size": 512}}) | |
| self.extract(state, texture_size=2048) | |
| self.assertEqual(state[app.GLB_SETTINGS_KEY], {"simplify": 0.95, "texture_size": 2048}) | |
| self.assertEqual(state[app.GLB_PAYLOAD_KEY], PAYLOAD) | |
| self.assertNotIn(app.GLB_MESH_KEY, state) | |
| def test_preparation_failure_is_surfaced_verbatim(self): | |
| state = cached_state(**{app.GLB_PAYLOAD_KEY: None}) | |
| with mock.patch.object(app, "_prepare_glb", | |
| side_effect=gr.Error("texturing failed")), \ | |
| mock.patch.object(app, "save_export_payload") as save: | |
| status, payload = app.extract_model(state, 0.95, app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| self.assertIsNone(payload) | |
| self.assertEqual(status["value"], "texturing failed") | |
| self.assertTrue(status["visible"]) | |
| save.assert_not_called() | |
| self.assertNotIn(app.GLB_MESH_KEY, state) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, ["gpu_task:bake"]) | |
| def test_a_failed_bake_is_retried_and_the_stale_error_clears(self): | |
| state = cached_state(**{app.GLB_PAYLOAD_KEY: None}) | |
| with mock.patch.object(app, "_prepare_glb", side_effect=RuntimeError("out of memory")): | |
| app.extract_model(state, 0.95, app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| self.assertEqual(state[app.GLB_ERROR_KEY], app.GLB_PREP_FAILURE_MESSAGE) | |
| (status, payload), _prepare = self.extract(state) | |
| self.assertEqual(payload, PAYLOAD) | |
| self.assertNotIn(app.GLB_ERROR_KEY, state) | |
| self.assertNotIn("visible", status) | |
| def test_packaging_failure_after_a_successful_bake_reports_the_packaging_error(self): | |
| state = cached_state(**{app.GLB_PAYLOAD_KEY: None}) | |
| error = gr.Error("Model was exported but failed validation for viewer rendering.") | |
| with mock.patch.object(app, "_prepare_glb", return_value={"glb": "baked"}), \ | |
| mock.patch.object(app, "save_export_payload", side_effect=error): | |
| status, payload = app.extract_model(state, 0.95, app.GLB_EXPORT_TEXTURE_SIZE, None) | |
| self.assertIsNone(payload) | |
| self.assertTrue(status["visible"]) | |
| self.assertIsNone(state.get(app.GLB_PAYLOAD_KEY)) | |
| def test_export_fallback_serves_the_cached_payload_without_allocation(self): | |
| entry = mock.Mock(args=(cached_state(), 0.95, app.GLB_EXPORT_TEXTURE_SIZE)) | |
| with mock.patch.object(app, "_prepare_glb") as prepare: | |
| self.assertEqual(app.run_pro_export(entry, None), PAYLOAD) | |
| prepare.assert_not_called() | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, []) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, []) | |
| def test_export_fallback_bakes_an_uncached_payload(self): | |
| entry = mock.Mock(args=({"runtime": "original"}, 0.95, app.GLB_EXPORT_TEXTURE_SIZE)) | |
| with mock.patch.object(app, "_prepare_glb", return_value={"glb": "baked"}) as prepare, \ | |
| mock.patch.object(app, "save_export_payload", return_value=dict(PAYLOAD)): | |
| self.assertEqual(app.run_pro_export(entry, None), PAYLOAD) | |
| self.assertEqual(prepare.call_args.args[1:], (app.GLB_EXPORT_TEXTURE_SIZE,)) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, ["gpu_task:bake"]) | |
| def test_export_fallback_raises_when_the_bake_fails(self): | |
| entry = mock.Mock(args=({"runtime": "original"}, 0.95, app.GLB_EXPORT_TEXTURE_SIZE)) | |
| with mock.patch.object(app, "_prepare_glb", side_effect=RuntimeError("out of memory")): | |
| with self.assertRaises(gr.Error): | |
| app.run_pro_export(entry, None) | |
| QUOTA_MESSAGE = "You have exceeded your ZeroGPU runs limit. Try again in 1 hour." | |
| class TextureSizeSelection(unittest.TestCase): | |
| def setUp(self): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| fallback_stubs.GPU_GATE["error"] = None | |
| def export(self, texture_size): | |
| state = {"runtime": "original"} | |
| with mock.patch.object(app, "_prepare_glb", return_value="prepared") as prepare, \ | |
| mock.patch.object(app, "save_export_payload", return_value=dict(PAYLOAD)): | |
| app.extract_model(state, app.GLB_EXPORT_SIMPLIFY, texture_size, None) | |
| return state, prepare | |
| def test_default_size_and_options_are_pinned(self): | |
| self.assertEqual(app.GLB_EXPORT_TEXTURE_SIZE, 1024) | |
| self.assertIn(512, app.GLB_TEXTURE_SIZES) | |
| self.assertEqual(tuple(app.GLB_TEXTURE_SIZES), (512, 1024, 1536, 2048, 4096)) | |
| def test_every_allowed_size_reaches_the_converter_and_is_recorded(self): | |
| for size in app.GLB_TEXTURE_SIZES: | |
| with self.subTest(size=size): | |
| state, prepare = self.export(size) | |
| prepare.assert_called_once() | |
| self.assertEqual(prepare.call_args.args[1:], (size,)) | |
| self.assertIs(prepare.call_args.args[0], state) | |
| self.assertEqual(state[app.GLB_SETTINGS_KEY], | |
| {"simplify": app.GLB_EXPORT_SIMPLIFY, "texture_size": size}) | |
| def test_to_glb_receives_the_chosen_size(self): | |
| from types import SimpleNamespace | |
| from export_profiles import simplify_ratio_for_texture_size | |
| source_faces = 400000 | |
| mesh = SimpleNamespace(faces=[0] * source_faces) | |
| gaussian = object() | |
| for size in app.GLB_TEXTURE_SIZES: | |
| with self.subTest(size=size): | |
| state = {"runtime": "original", "mesh": {}, "gaussian": {}} | |
| with mock.patch.object(app, "HAS_DIFF_GAUSSIAN_RASTERIZATION", True), \ | |
| mock.patch.object(app, "unpack_original_state", | |
| return_value={"mesh": mesh, "gaussian": gaussian}), \ | |
| mock.patch.object(app.trellis_postprocessing_utils, "to_glb", | |
| return_value="prepared") as convert, \ | |
| mock.patch.object(app, "save_export_payload", return_value=dict(PAYLOAD)): | |
| app.extract_model(state, app.GLB_EXPORT_SIMPLIFY, size, None) | |
| self.assertEqual(convert.call_args.kwargs["texture_size"], size) | |
| self.assertEqual(convert.call_args.kwargs["simplify"], | |
| simplify_ratio_for_texture_size(size, source_faces)) | |
| self.assertEqual(state[app.GLB_SETTINGS_KEY]["texture_size"], size) | |
| def test_public_generation_defaults_to_the_export_size(self): | |
| expected = app.GLB_EXPORT_TEXTURE_SIZE | |
| for function in (app.generate_preview_frames, app.generate_and_extract, | |
| app._generate_and_render_model): | |
| with self.subTest(function=function.__name__): | |
| params = inspect.signature(function).parameters | |
| self.assertIn("texture_size", params) | |
| self.assertEqual(params["texture_size"].default, expected) | |
| self.assertEqual(tuple(inspect.signature(app.generate_and_extract).parameters)[:7], | |
| ("image", "seed", "ss_guidance_strength", "ss_sampling_steps", | |
| "slat_guidance_strength", "slat_sampling_steps", "texture_size")) | |
| def test_invalid_sizes_are_refused(self): | |
| for size in app.GLB_TEXTURE_SIZES: | |
| self.assertEqual(app.validate_texture_size(size), size) | |
| for value in (768, 0, -512, None, ""): | |
| with self.subTest(value=value): | |
| with self.assertRaises(gr.Error): | |
| app.validate_texture_size(value) | |
| def test_exports_publish_the_recorded_size(self): | |
| state = {app.GLB_SETTINGS_KEY: {"simplify": app.GLB_EXPORT_SIMPLIFY, "texture_size": 1536}} | |
| self.assertEqual(app.enable_model_exports(state)[3], 1536) | |
| self.assertEqual(app.enable_model_exports({"runtime": "original"})[3], app.GLB_EXPORT_TEXTURE_SIZE) | |
| self.assertEqual(app.enable_model_exports(None)[3], app.GLB_EXPORT_TEXTURE_SIZE) | |
| class GenerationTextureSize(unittest.TestCase): | |
| def setUp(self): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| fallback_stubs.GPU_GATE["error"] = None | |
| self.prepared = [] | |
| for patcher in ( | |
| mock.patch.object(app, "pipeline", object()), | |
| mock.patch.object(app, "preprocess_image", side_effect=lambda image: image), | |
| mock.patch.object(app, "_run_original_pipeline", | |
| return_value={"mesh": [0], "gaussian": ["gaussian-0"]}), | |
| mock.patch.object(app, "pack_original_state", | |
| return_value={"runtime": "original", "marker": "state-0"}), | |
| mock.patch.object(app, "render_original_preview_frames", return_value=[{"frame": 0}]), | |
| mock.patch.object(app, "_prepare_glb", side_effect=self.prepare), | |
| mock.patch.object(app, "save_export_payload", return_value=dict(PAYLOAD)), | |
| mock.patch.object(app, "write_mp4", return_value=True), | |
| ): | |
| patcher.start() | |
| self.addCleanup(patcher.stop) | |
| def prepare(self, state, texture_size): | |
| self.prepared.append(texture_size) | |
| return {"glb": "prepared"} | |
| def generate(self, texture_size): | |
| image = Image.new("RGB", (8, 8), (12, 34, 56)) | |
| return app.generate_and_extract(image, 0, 6.0, 12, 1.0, 12, texture_size) | |
| def test_generation_returns_a_preview_without_baking_a_texture(self): | |
| for size in app.GLB_TEXTURE_SIZES: | |
| with self.subTest(size=size): | |
| fallback_stubs.GPU_CALLS.clear() | |
| self.prepared.clear() | |
| frames, state = self.generate(size) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, ["gpu_task:preview"]) | |
| self.assertEqual(self.prepared, []) | |
| self.assertTrue(frames.endswith(".mp4"), frames) | |
| self.assertEqual(state[app.GLB_REQUESTED_KEY], | |
| {"simplify": app.GLB_EXPORT_SIMPLIFY, "texture_size": size}) | |
| self.assertNotIn(app.GLB_PAYLOAD_KEY, state) | |
| def test_a_matching_export_bakes_once_and_is_then_served_from_cache(self): | |
| for size in app.GLB_TEXTURE_SIZES: | |
| with self.subTest(size=size): | |
| fallback_stubs.GPU_CALLS.clear() | |
| self.prepared.clear() | |
| _frames, state = self.generate(size) | |
| status, payload = app.extract_model(state, app.GLB_EXPORT_SIMPLIFY, size, None) | |
| self.assertEqual(payload, PAYLOAD) | |
| self.assertEqual(self.prepared, [size]) | |
| self.assertEqual(state[app.GLB_SETTINGS_KEY], | |
| {"simplify": app.GLB_EXPORT_SIMPLIFY, "texture_size": size}) | |
| self.assertIn("Finalizing", status["value"]) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, | |
| ["gpu_task:preview", "gpu_task:bake"]) | |
| self.assertEqual(app.cached_glb_payload(state, app.GLB_EXPORT_SIMPLIFY, size), PAYLOAD) | |
| _status, again = app.extract_model(state, app.GLB_EXPORT_SIMPLIFY, size, None) | |
| self.assertEqual(again, PAYLOAD) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, | |
| ["gpu_task:preview", "gpu_task:bake"]) | |
| def test_invalid_size_is_rejected_before_any_gpu_attempt(self): | |
| for value in (768, 0, None): | |
| with self.subTest(value=value): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| self.prepared.clear() | |
| with self.assertRaises(gr.Error): | |
| self.generate(value) | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, []) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, []) | |
| self.assertEqual(self.prepared, []) | |
| def test_quota_denial_defers_the_chosen_size_and_replay_keeps_it(self): | |
| size = 1536 | |
| fallback_stubs.GPU_GATE["error"] = gr.Error(QUOTA_MESSAGE) | |
| deferred = [] | |
| with mock.patch.object(app, "defer_pending_fallback", | |
| side_effect=lambda op, args, req, nonce=None: deferred.append((op, args)) or True): | |
| frames, state = self.generate(size) | |
| self.assertIsNone(frames) | |
| self.assertIsNone(state) | |
| operation, args = deferred[0] | |
| self.assertEqual(operation, app.GENERATE) | |
| self.assertEqual(args[-1], size) | |
| self.assertEqual(tuple(args[1:6]), (0, 6.0, 12, 1.0, 12)) | |
| self.assertIsInstance(args[0], Image.Image) | |
| fallback_stubs.GPU_GATE["error"] = None | |
| self.prepared.clear() | |
| _video, replayed = app.run_pro_generate(mock.Mock(args=args), None) | |
| deadline = time.time() + 5 | |
| while app.GLB_PAYLOAD_KEY not in replayed and time.time() < deadline: | |
| time.sleep(0.01) | |
| self.assertEqual(self.prepared, [size]) | |
| self.assertEqual(replayed[app.GLB_REQUESTED_KEY]["texture_size"], size) | |
| self.assertEqual(app.extract_model(replayed, app.GLB_EXPORT_SIMPLIFY, size, None)[1], PAYLOAD) | |
| self.assertEqual(self.prepared, [size]) | |
| self.assertEqual(app.enable_model_exports(replayed)[3], size) | |
| def test_cached_export_follows_the_generated_size_after_a_later_choice(self): | |
| _frames, state = self.generate(512) | |
| attempts_before = list(fallback_stubs.GPU_ATTEMPTS) | |
| self.assertEqual(attempts_before, ["gpu_task:preview"]) | |
| _status, payload = app.extract_model(state, app.GLB_EXPORT_SIMPLIFY, 512, None) | |
| self.assertEqual(payload, PAYLOAD) | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, | |
| ["gpu_task:preview", "gpu_task:bake"]) | |
| self.prepared.clear() | |
| _status, payload = app.extract_model(state, app.GLB_EXPORT_SIMPLIFY, 1024, None) | |
| self.assertEqual(payload, PAYLOAD) | |
| self.assertEqual(self.prepared, [1024]) | |
| self.assertEqual(state[app.GLB_SETTINGS_KEY]["texture_size"], 1024) | |
| self.assertEqual(app.cached_glb_payload(state, app.GLB_EXPORT_SIMPLIFY, 1024), PAYLOAD) | |
| self.assertEqual(app.enable_model_exports(state)[3], 512) | |
| self.assertEqual(state[app.GLB_REQUESTED_KEY]["texture_size"], 512) | |
| class RealPayloadPipeline(unittest.TestCase): | |
| def setUp(self): | |
| fallback_stubs.GPU_ATTEMPTS.clear() | |
| fallback_stubs.GPU_CALLS.clear() | |
| fallback_stubs.GPU_GATE["error"] = None | |
| self.export_dir = tempfile.TemporaryDirectory() | |
| self.addCleanup(self.export_dir.cleanup) | |
| for patcher in (mock.patch.object(app, "EXPORT_DIR", self.export_dir.name), | |
| mock.patch.object(app, "RUNNING_ON_SPACES", False)): | |
| patcher.start() | |
| self.addCleanup(patcher.stop) | |
| def generated_state(self): | |
| return { | |
| "runtime": "original", | |
| app.GLB_MESH_KEY: self.textured_mesh(), | |
| app.GLB_SETTINGS_KEY: dict(SETTINGS), | |
| } | |
| def textured_mesh(self): | |
| vertices = np.random.RandomState(0).rand(64, 3).astype(np.float32) | |
| faces = np.arange(60, dtype=np.int64).reshape(20, 3) | |
| uv = np.random.RandomState(1).rand(64, 2).astype(np.float32) | |
| texture = Image.fromarray((np.random.RandomState(2).rand(64, 64, 4) * 255).astype(np.uint8), mode="RGBA") | |
| material = trimesh.visual.material.PBRMaterial(baseColorTexture=texture) | |
| return trimesh.Trimesh(vertices=vertices, faces=faces, process=False, | |
| visual=trimesh.visual.TextureVisuals(uv=uv, material=material)) | |
| def test_prepared_glb_is_exported_validated_and_viewable(self): | |
| state = self.generated_state() | |
| app.cache_prepared_glb(state) | |
| self.assertNotIn(app.GLB_MESH_KEY, state) | |
| self.assertNotIn(app.GLB_ERROR_KEY, state) | |
| payload = state[app.GLB_PAYLOAD_KEY] | |
| self.assertTrue(payload["glb_name"].endswith(".glb")) | |
| status, returned = app.extract_model(state, 0.95, app.GLB_EXPORT_TEXTURE_SIZE) | |
| self.assertEqual(returned, payload) | |
| self.assertIn("Finalizing", status["value"]) | |
| update, glb_path = app.materialize_export_payload(returned, None) | |
| self.assertTrue(os.path.isfile(glb_path)) | |
| self.assertGreater(os.path.getsize(glb_path), 0) | |
| reloaded = trimesh.load(glb_path, force="scene") | |
| self.assertTrue(reloaded.geometry) | |
| self.assertEqual(fallback_stubs.GPU_ATTEMPTS, []) | |
| self.assertEqual(fallback_stubs.GPU_CALLS, []) | |
| def test_payload_survives_serialization(self): | |
| state = self.generated_state() | |
| app.cache_prepared_glb(state) | |
| restored = pickle.loads(pickle.dumps(state)) | |
| self.assertEqual(restored[app.GLB_PAYLOAD_KEY], state[app.GLB_PAYLOAD_KEY]) | |
| self.assertEqual(app.extract_model(restored, 0.95, app.GLB_EXPORT_TEXTURE_SIZE)[1], | |
| state[app.GLB_PAYLOAD_KEY]) | |
| if __name__ == "__main__": | |
| unittest.main() | |