diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index 34c6143ef..68a185c40 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -1,3 +1,4 @@ +from av.bitstream import BitStreamFilterContext from av.container import InputContainer from av.subtitles.stream import SubtitleStream from av.video.reformatter import ColorPrimaries, ColorRange, ColorTrc @@ -96,6 +97,31 @@ def video_stream_bit_depth(stream) -> int: return max(component.bits for component in stream.format.components) +def isobmff_hevc_filter(output_container, stream, out_stream): + """Apple players need the 'hvc1' sample entry, not FFmpeg's default 'hev1'. Annex B input without + extradata makes the muxer build hvcC from the first packet and strip in-band parameter sets; + 'hvc1' sources already have a complete hvcC and only need the tag PyAV reset.""" + if output_container.format.name not in ("mp4", "mov") or stream.codec.canonical_name != "hevc": + return None + try: + codec_tag = stream.codec_context.codec_tag + except UnicodeDecodeError: + codec_tag = "" + if codec_tag == "hvc1": + out_stream.codec_context.codec_tag = "hvc1" + return None + hevc_filter = BitStreamFilterContext("hevc_mp4toannexb", stream, out_stream) + out_stream.codec_context.codec_tag = "hvc1" + out_stream.codec_context.extradata = None + return hevc_filter + + +def filter_hevc_packet(hevc_filter, packet): + if packet.has_sidedata("new_extradata"): + raise ValueError("HEVC with multiple sample descriptions cannot be remuxed as hvc1; re-encode it instead") + return hevc_filter.filter(packet) + + def last_decodable_audio_stream(container: InputContainer): """Streams FFmpeg has no decoder for have no codec context, and decoding their packets crashes the process (e.g. APAC spatial-audio track in iPhone).""" @@ -601,19 +627,26 @@ class VideoFromFile(VideoInput): # Add streams to the new container. Streams with no codec context cannot be used as an output template. stream_map = {} + hevc_filters = {} for stream in streams: if isinstance(stream, (av.VideoStream, av.AudioStream, SubtitleStream)): if stream.codec_context is None: logging.warning("Skipping %s stream %d with unsupported codec", stream.type, stream.index) continue out_stream = output_container.add_stream_from_template(template=stream, opaque=True) + hevc_filter = isobmff_hevc_filter(output_container, stream, out_stream) + if hevc_filter is not None: + hevc_filters[stream] = hevc_filter stream_map[stream] = out_stream # Write packets to the new container for packet in container.demux(): if packet.stream in stream_map and packet.dts is not None: - packet.stream = stream_map[packet.stream] - output_container.mux(packet) + out_stream = stream_map[packet.stream] + hevc_filter = hevc_filters.get(packet.stream) + for out_packet in filter_hevc_packet(hevc_filter, packet) if hevc_filter else (packet,): + out_packet.stream = out_stream + output_container.mux(out_packet) def _save_transcoded( self, diff --git a/tests-unit/comfy_api_test/video_types_test.py b/tests-unit/comfy_api_test/video_types_test.py index fe7b8de2a..6b50b86d5 100644 --- a/tests-unit/comfy_api_test/video_types_test.py +++ b/tests-unit/comfy_api_test/video_types_test.py @@ -1289,6 +1289,171 @@ def test_save_to_transcode_bakes_rotation(): os.unlink(file_path) +def hevc_encoder_available(): + try: + av.Codec("libx265", "w") + return True + except av.codec.codec.UnknownCodecError: + return False + + +hevc_remux_test = pytest.mark.skipif(not hevc_encoder_available(), reason="libx265 encoder not available") + + +def create_hevc_mp4(codec_tag=None, x265_params=None): + """In-memory HEVC mp4, which FFmpeg tags 'hev1' unless codec_tag is given.""" + buffer = io.BytesIO() + options = {"x265-params": ":".join(["log-level=none"] + ([x265_params] if x265_params else []))} + with av.open(buffer, mode="w", format="mp4") as container: + stream = container.add_stream("libx265", rate=30, options=options) + stream.width = 64 + stream.height = 64 + stream.pix_fmt = "yuv420p" + if codec_tag is not None: + stream.codec_context.codec_tag = codec_tag + for i in range(3): + frame = av.VideoFrame.from_ndarray( + torch.ones(64, 64, 3, dtype=torch.uint8).numpy() * (i * 85), + format="rgb24", + ).reformat(format="yuv420p") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + buffer.seek(0) + return buffer + + +def filter_hvcc_arrays(source: io.BytesIO, keep_nal_types: tuple = ()) -> io.BytesIO: + """hev1 whose hvcC keeps only the given NAL array types (32=VPS, 33=SPS, 34=PPS, 39=SEI).""" + output = io.BytesIO() + with av.open(source) as container, av.open(output, mode="w", format="mp4") as output_container: + stream = container.streams.video[0] + extradata = stream.codec_context.extradata + rebuilt, pos = bytearray(extradata[:22] + b"\x00"), 23 + for _ in range(extradata[22]): + array_start = pos + num_nalus = int.from_bytes(extradata[pos + 1:pos + 3], "big") + pos += 3 + for _ in range(num_nalus): + pos += 2 + int.from_bytes(extradata[pos:pos + 2], "big") + if extradata[array_start] & 0x3F in keep_nal_types: + rebuilt += extradata[array_start:pos] + rebuilt[22] += 1 + out_stream = output_container.add_stream_from_template(template=stream, opaque=True) + out_stream.codec_context.extradata = bytes(rebuilt) + for packet in container.demux(stream): + if packet.dts is not None: + packet.stream = out_stream + output_container.mux(packet) + output.seek(0) + return output + + +def concat_hevc_sample_descriptions(first: io.BytesIO, second: io.BytesIO, codec_tag=None) -> io.BytesIO: + """mp4 with two sample descriptions, the second created from new_extradata packet side data.""" + output = io.BytesIO() + with av.open(first) as a, av.open(second) as b, av.open(output, mode="w", format="mp4") as output_container: + out_stream = output_container.add_stream_from_template(template=a.streams.video[0], opaque=True) + if codec_tag is not None: + out_stream.codec_context.codec_tag = codec_tag + end = 0 + for packet in a.demux(a.streams.video[0]): + if packet.dts is not None: + packet.stream = out_stream + output_container.mux(packet) + end = max(end, packet.pts + packet.duration) + extradata = b.streams.video[0].codec_context.extradata + packets = [packet for packet in b.demux(b.streams.video[0]) if packet.dts is not None] + sidedata = av.packet.PacketSideData(av.packet.packet_sidedata_type_from_literal("new_extradata"), len(extradata)) + sidedata.update(extradata) + packets[0].set_sidedata(sidedata) + for packet in packets: + packet.pts += end + packet.dts += end + packet.stream = out_stream + output_container.mux(packet) + output.seek(0) + return output + + +def sample_description_count(data: bytes, start: int = 0, end: int | None = None) -> int: + """Number of 'stsd' entries; a broken orphan entry is invisible to decoders but not to players.""" + pos, end, count = start, len(data) if end is None else end, 0 + while pos + 8 <= end: + size, box = int.from_bytes(data[pos:pos + 4], "big"), data[pos + 4:pos + 8] + if box == b"stsd": + count += int.from_bytes(data[pos + 12:pos + 16], "big") + elif box in (b"moov", b"trak", b"mdia", b"minf", b"stbl"): + count += sample_description_count(data, pos + 8, pos + size) + pos += size + return count + + +def probe_hevc(source: io.BytesIO) -> dict: + source.seek(0) + with av.open(source) as container: + stream = container.streams.video[0] + return { + "tag": stream.codec_context.codec_tag, + "frames": sum(1 for packet in container.demux(stream) for _ in packet.decode()), + "sample_descriptions": sample_description_count(source.getvalue()), + } + + +def remux_and_probe(source: io.BytesIO, **save_kwargs) -> dict: + output = io.BytesIO() + source.seek(0) + VideoFromFile(source).save_to(output, **save_kwargs) + return probe_hevc(output) + + +@hevc_remux_test +@pytest.mark.parametrize("save_kwargs", [pytest.param({}, id="mov"), pytest.param({"format": VideoContainer.MP4}, id="mp4")]) +@pytest.mark.parametrize("codec_tag", [None, "dvh1"], ids=["hev1", "dvh1"]) +def test_save_to_remux_retags_hevc_as_hvc1(codec_tag, save_kwargs): + """Remuxed HEVC gets 'hvc1' instead of FFmpeg's default 'hev1', Dolby Vision included: PyAV + resets the source tag and the muxer drops the DV boxes anyway, so skipping 'dvh1' means 'hev1'.""" + source = create_hevc_mp4(codec_tag=codec_tag) + assert probe_hevc(source)["tag"] == (codec_tag or "hev1") + assert remux_and_probe(source, **save_kwargs) == {"tag": "hvc1", "frames": 3, "sample_descriptions": 1} + + +@hevc_remux_test +@pytest.mark.parametrize( + "keep_nal_types", + [pytest.param((), id="empty-hvcc"), pytest.param((33, 34), id="missing-vps"), pytest.param((39,), id="sei-only")], +) +def test_save_to_remux_rebuilds_hvcc_from_inband_parameter_sets(keep_nal_types): + """Parameter sets missing from hvcC but present in-band yield one valid hvcC, not an empty one.""" + source = filter_hvcc_arrays(create_hevc_mp4(x265_params="repeat-headers=1"), keep_nal_types) + output = io.BytesIO() + VideoFromFile(source).save_to(output) + output.seek(0) + with av.open(output) as container: + assert container.streams.video[0].codec_context.extradata[22] > 0 # hvcC numOfArrays + assert probe_hevc(output) == {"tag": "hvc1", "frames": 3, "sample_descriptions": 1} + + +@hevc_remux_test +def test_save_to_remux_keeps_hvc1_sources_untouched(): + """'hvc1' sources skip the bitstream filter, so their sample descriptions survive as they are.""" + source = concat_hevc_sample_descriptions( + create_hevc_mp4(codec_tag="hvc1"), create_hevc_mp4(codec_tag="hvc1", x265_params="no-sao=1"), codec_tag="hvc1" + ) + assert probe_hevc(source) == {"tag": "hvc1", "frames": 6, "sample_descriptions": 2} + assert remux_and_probe(source) == {"tag": "hvc1", "frames": 6, "sample_descriptions": 2} + + +@hevc_remux_test +def test_save_to_remux_rejects_hevc_with_multiple_sample_descriptions(): + """hevc_mp4toannexb ignores mid-stream parameter set changes, so such sources must be re-encoded.""" + source = concat_hevc_sample_descriptions(create_hevc_mp4(), create_hevc_mp4(x265_params="no-sao=1")) + assert probe_hevc(source) == {"tag": "hev1", "frames": 6, "sample_descriptions": 2} + with pytest.raises(ValueError, match="multiple sample descriptions"): + remux_and_probe(source) + + def test_save_to_transcode_skips_undecodable_audio(): """Streaming transcode keeps the decodable audio track and drops undecodable ones; with no decodable audio at all the output is video-only instead of crashing."""