Skip to content

Commit 00fff60

Browse files
guillKosinkadink
andauthored
feat(jobs): add 3d to PREVIEWABLE_MEDIA_TYPES for first-class 3D output support (Comfy-Org#12381)
Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>
1 parent 123a787 commit 00fff60

File tree

2 files changed

+269
-14
lines changed

2 files changed

+269
-14
lines changed

comfy_execution/jobs.py

Lines changed: 65 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,60 @@ class JobStatus:
2020

2121

2222
# Media types that can be previewed in the frontend
23-
PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio'})
23+
PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d'})
2424

2525
# 3D file extensions for preview fallback (no dedicated media_type exists)
26-
THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb'})
26+
THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb', '.usdz'})
27+
28+
29+
def has_3d_extension(filename: str) -> bool:
30+
lower = filename.lower()
31+
return any(lower.endswith(ext) for ext in THREE_D_EXTENSIONS)
32+
33+
34+
def normalize_output_item(item):
35+
"""Normalize a single output list item for the jobs API.
36+
37+
Returns the normalized item, or None to exclude it.
38+
String items with 3D extensions become {filename, type, subfolder} dicts.
39+
"""
40+
if item is None:
41+
return None
42+
if isinstance(item, str):
43+
if has_3d_extension(item):
44+
return {'filename': item, 'type': 'output', 'subfolder': '', 'mediaType': '3d'}
45+
return None
46+
if isinstance(item, dict):
47+
return item
48+
return None
49+
50+
51+
def normalize_outputs(outputs: dict) -> dict:
52+
"""Normalize raw node outputs for the jobs API.
53+
54+
Transforms string 3D filenames into file output dicts and removes
55+
None items. All other items (non-3D strings, dicts, etc.) are
56+
preserved as-is.
57+
"""
58+
normalized = {}
59+
for node_id, node_outputs in outputs.items():
60+
if not isinstance(node_outputs, dict):
61+
normalized[node_id] = node_outputs
62+
continue
63+
normalized_node = {}
64+
for media_type, items in node_outputs.items():
65+
if media_type == 'animated' or not isinstance(items, list):
66+
normalized_node[media_type] = items
67+
continue
68+
normalized_items = []
69+
for item in items:
70+
if item is None:
71+
continue
72+
norm = normalize_output_item(item)
73+
normalized_items.append(norm if norm is not None else item)
74+
normalized_node[media_type] = normalized_items
75+
normalized[node_id] = normalized_node
76+
return normalized
2777

2878

2979
def _extract_job_metadata(extra_data: dict) -> tuple[Optional[int], Optional[str]]:
@@ -45,9 +95,9 @@ def is_previewable(media_type: str, item: dict) -> bool:
4595
Maintains backwards compatibility with existing logic.
4696
4797
Priority:
48-
1. media_type is 'images', 'video', or 'audio'
98+
1. media_type is 'images', 'video', 'audio', or '3d'
4999
2. format field starts with 'video/' or 'audio/'
50-
3. filename has a 3D extension (.obj, .fbx, .gltf, .glb)
100+
3. filename has a 3D extension (.obj, .fbx, .gltf, .glb, .usdz)
51101
"""
52102
if media_type in PREVIEWABLE_MEDIA_TYPES:
53103
return True
@@ -139,7 +189,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
139189
})
140190

141191
if include_outputs:
142-
job['outputs'] = outputs
192+
job['outputs'] = normalize_outputs(outputs)
143193
job['execution_status'] = status_info
144194
job['workflow'] = {
145195
'prompt': prompt,
@@ -171,18 +221,23 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
171221
continue
172222

173223
for item in items:
224+
normalized = normalize_output_item(item)
225+
if normalized is None:
226+
continue
227+
174228
count += 1
175229

176-
if not isinstance(item, dict):
230+
if preview_output is not None:
177231
continue
178232

179-
if preview_output is None and is_previewable(media_type, item):
233+
if isinstance(normalized, dict) and is_previewable(media_type, normalized):
180234
enriched = {
181-
**item,
235+
**normalized,
182236
'nodeId': node_id,
183-
'mediaType': media_type
184237
}
185-
if item.get('type') == 'output':
238+
if 'mediaType' not in normalized:
239+
enriched['mediaType'] = media_type
240+
if normalized.get('type') == 'output':
186241
preview_output = enriched
187242
elif fallback_preview is None:
188243
fallback_preview = enriched

tests/execution/test_jobs.py

Lines changed: 204 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,11 @@
55
is_previewable,
66
normalize_queue_item,
77
normalize_history_item,
8+
normalize_output_item,
9+
normalize_outputs,
810
get_outputs_summary,
911
apply_sorting,
12+
has_3d_extension,
1013
)
1114

1215

@@ -35,8 +38,8 @@ class TestIsPreviewable:
3538
"""Unit tests for is_previewable()"""
3639

3740
def test_previewable_media_types(self):
38-
"""Images, video, audio media types should be previewable."""
39-
for media_type in ['images', 'video', 'audio']:
41+
"""Images, video, audio, 3d media types should be previewable."""
42+
for media_type in ['images', 'video', 'audio', '3d']:
4043
assert is_previewable(media_type, {}) is True
4144

4245
def test_non_previewable_media_types(self):
@@ -46,7 +49,7 @@ def test_non_previewable_media_types(self):
4649

4750
def test_3d_extensions_previewable(self):
4851
"""3D file extensions should be previewable regardless of media_type."""
49-
for ext in ['.obj', '.fbx', '.gltf', '.glb']:
52+
for ext in ['.obj', '.fbx', '.gltf', '.glb', '.usdz']:
5053
item = {'filename': f'model{ext}'}
5154
assert is_previewable('files', item) is True
5255

@@ -160,7 +163,7 @@ def test_previewable_media_types(self):
160163

161164
def test_3d_files_previewable(self):
162165
"""3D file extensions should be previewable."""
163-
for ext in ['.obj', '.fbx', '.gltf', '.glb']:
166+
for ext in ['.obj', '.fbx', '.gltf', '.glb', '.usdz']:
164167
outputs = {
165168
'node1': {
166169
'files': [{'filename': f'model{ext}', 'type': 'output'}]
@@ -192,6 +195,64 @@ def test_preview_enriched_with_node_metadata(self):
192195
assert preview['mediaType'] == 'images'
193196
assert preview['subfolder'] == 'outputs'
194197

198+
def test_string_3d_filename_creates_preview(self):
199+
"""String items with 3D extensions should synthesize a preview (Preview3D node output).
200+
Only the .glb counts — nulls and non-file strings are excluded."""
201+
outputs = {
202+
'node1': {
203+
'result': ['preview3d_abc123.glb', None, None]
204+
}
205+
}
206+
count, preview = get_outputs_summary(outputs)
207+
assert count == 1
208+
assert preview is not None
209+
assert preview['filename'] == 'preview3d_abc123.glb'
210+
assert preview['mediaType'] == '3d'
211+
assert preview['nodeId'] == 'node1'
212+
assert preview['type'] == 'output'
213+
214+
def test_string_non_3d_filename_no_preview(self):
215+
"""String items without 3D extensions should not create a preview."""
216+
outputs = {
217+
'node1': {
218+
'result': ['data.json', None]
219+
}
220+
}
221+
count, preview = get_outputs_summary(outputs)
222+
assert count == 0
223+
assert preview is None
224+
225+
def test_string_3d_filename_used_as_fallback(self):
226+
"""String 3D preview should be used when no dict items are previewable."""
227+
outputs = {
228+
'node1': {
229+
'latents': [{'filename': 'latent.safetensors'}],
230+
},
231+
'node2': {
232+
'result': ['model.glb', None]
233+
}
234+
}
235+
count, preview = get_outputs_summary(outputs)
236+
assert preview is not None
237+
assert preview['filename'] == 'model.glb'
238+
assert preview['mediaType'] == '3d'
239+
240+
241+
class TestHas3DExtension:
242+
"""Unit tests for has_3d_extension()"""
243+
244+
def test_recognized_extensions(self):
245+
for ext in ['.obj', '.fbx', '.gltf', '.glb', '.usdz']:
246+
assert has_3d_extension(f'model{ext}') is True
247+
248+
def test_case_insensitive(self):
249+
assert has_3d_extension('MODEL.GLB') is True
250+
assert has_3d_extension('Scene.GLTF') is True
251+
252+
def test_non_3d_extensions(self):
253+
for name in ['photo.png', 'video.mp4', 'data.json', 'model']:
254+
assert has_3d_extension(name) is False
255+
195256

196257
class TestApplySorting:
197258
"""Unit tests for apply_sorting()"""
@@ -395,3 +456,142 @@ def test_include_outputs(self):
395456
'prompt': {'nodes': {'1': {}}},
396457
'extra_data': {'create_time': 1234567890, 'client_id': 'abc'},
397458
}
459+
460+
def test_include_outputs_normalizes_3d_strings(self):
461+
"""Detail view should transform string 3D filenames into file output dicts."""
462+
history_item = {
463+
'prompt': (
464+
5,
465+
'prompt-3d',
466+
{'nodes': {}},
467+
{'create_time': 1234567890},
468+
['node1'],
469+
),
470+
'status': {'status_str': 'success', 'completed': True, 'messages': []},
471+
'outputs': {
472+
'node1': {
473+
'result': ['preview3d_abc123.glb', None, None]
474+
}
475+
},
476+
}
477+
job = normalize_history_item('prompt-3d', history_item, include_outputs=True)
478+
479+
assert job['outputs_count'] == 1
480+
result_items = job['outputs']['node1']['result']
481+
assert len(result_items) == 1
482+
assert result_items[0] == {
483+
'filename': 'preview3d_abc123.glb',
484+
'type': 'output',
485+
'subfolder': '',
486+
'mediaType': '3d',
487+
}
488+
489+
def test_include_outputs_preserves_dict_items(self):
490+
"""Detail view normalization should pass dict items through unchanged."""
491+
history_item = {
492+
'prompt': (
493+
5,
494+
'prompt-img',
495+
{'nodes': {}},
496+
{'create_time': 1234567890},
497+
['node1'],
498+
),
499+
'status': {'status_str': 'success', 'completed': True, 'messages': []},
500+
'outputs': {
501+
'node1': {
502+
'images': [
503+
{'filename': 'photo.png', 'type': 'output', 'subfolder': ''},
504+
]
505+
}
506+
},
507+
}
508+
job = normalize_history_item('prompt-img', history_item, include_outputs=True)
509+
510+
assert job['outputs_count'] == 1
511+
assert job['outputs']['node1']['images'] == [
512+
{'filename': 'photo.png', 'type': 'output', 'subfolder': ''},
513+
]
514+
515+
516+
class TestNormalizeOutputItem:
517+
"""Unit tests for normalize_output_item()"""
518+
519+
def test_none_returns_none(self):
520+
assert normalize_output_item(None) is None
521+
522+
def test_string_3d_extension_synthesizes_dict(self):
523+
result = normalize_output_item('model.glb')
524+
assert result == {'filename': 'model.glb', 'type': 'output', 'subfolder': '', 'mediaType': '3d'}
525+
526+
def test_string_non_3d_extension_returns_none(self):
527+
assert normalize_output_item('data.json') is None
528+
529+
def test_string_no_extension_returns_none(self):
530+
assert normalize_output_item('camera_info_string') is None
531+
532+
def test_dict_passes_through(self):
533+
item = {'filename': 'test.png', 'type': 'output'}
534+
assert normalize_output_item(item) is item
535+
536+
def test_other_types_return_none(self):
537+
assert normalize_output_item(42) is None
538+
assert normalize_output_item(True) is None
539+
540+
541+
class TestNormalizeOutputs:
542+
"""Unit tests for normalize_outputs()"""
543+
544+
def test_empty_outputs(self):
545+
assert normalize_outputs({}) == {}
546+
547+
def test_dict_items_pass_through(self):
548+
outputs = {
549+
'node1': {
550+
'images': [{'filename': 'a.png', 'type': 'output'}],
551+
}
552+
}
553+
result = normalize_outputs(outputs)
554+
assert result == outputs
555+
556+
def test_3d_string_synthesized(self):
557+
outputs = {
558+
'node1': {
559+
'result': ['model.glb', None, None],
560+
}
561+
}
562+
result = normalize_outputs(outputs)
563+
assert result == {
564+
'node1': {
565+
'result': [
566+
{'filename': 'model.glb', 'type': 'output', 'subfolder': '', 'mediaType': '3d'},
567+
],
568+
}
569+
}
570+
571+
def test_animated_key_preserved(self):
572+
outputs = {
573+
'node1': {
574+
'images': [{'filename': 'a.png', 'type': 'output'}],
575+
'animated': [True],
576+
}
577+
}
578+
result = normalize_outputs(outputs)
579+
assert result['node1']['animated'] == [True]
580+
581+
def test_non_dict_node_outputs_preserved(self):
582+
outputs = {'node1': 'unexpected_value'}
583+
result = normalize_outputs(outputs)
584+
assert result == {'node1': 'unexpected_value'}
585+
586+
def test_none_items_filtered_but_other_types_preserved(self):
587+
outputs = {
588+
'node1': {
589+
'result': ['data.json', None, [1, 2, 3]],
590+
}
591+
}
592+
result = normalize_outputs(outputs)
593+
assert result == {
594+
'node1': {
595+
'result': ['data.json', [1, 2, 3]],
596+
}
597+
}

0 commit comments

Comments
 (0)