FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix(tasks): #775 - finishing commands load an image file with its alpha · dkackman/diffusers-workflow@b1b07b1 · GitHub

Repository navigation

Commit b1b07b1

Browse files
andcommitted
fix(tasks): #775 - finishing commands load an image file with its alpha
grade, sharpen, film_grain and apply_lut read a media file through fetch_image, which diffusers' load_image flattened to RGB, so an RGBA asset lost its alpha before split_alpha ever saw it. fetch_image gains keep_alpha; _load_media sets it. Pipeline inputs still load RGB. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent c7578ae commit b1b07b1

4 files changed

Lines changed: 77 additions & 7 deletions

File tree

‎docs/ARCHITECTURE.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ one to open.
7474
| Join processed windows back to the source's length | `dw/tasks/windows.py`: `join_windows`; `dw/task_domains.py`: `window_count`, `window_count_problem`, `join_windows_errors` | Output is exactly the source's frames, each once: later windows' first `overlap` frames blend over the previous window's last on the open ramp `dissolve_videos` uses. The window count is `window_count`'s `ceil(source_frames / stride)` and nothing else (it is also `window_video`'s last index plus one). Audio is the source's own track; shots are one per window with `overlap_frames` on each seam, their sample side laid over that track by `dw/shots.py`'s `remeasured_shots` (the same rounding as `frames_to_samples`, so the spans are cumulative). | `tests/test_join_windows.py::TestRoundTrip::test_unprocessed_windows_rejoin_to_the_source`, `tests/test_join_windows.py::TestShots::test_spans_partition_the_frames_and_tile_the_audio`, `tests/test_join_windows.py::TestUint8Output::test_matches_the_float32_accumulator_exactly` |
7575
| A `join_windows` list of the wrong length, before the run | `dw/window_count_errors.py`: `window_count_errors`, wired into `validation_errors` | Owned by the task, never a template: any `join_windows` step whose `source` is probeable header-only (`asset:`/`output:` or a readable literal path), with literal `num_frames`/`overlap` and an expanded `videos` list, is checked against `window_count_problem` - the same rule the run refuses with. Anything not knowable stays silent for the run. | `tests/test_window_count_errors.py` |
7676
| Fit to a model and restore | `dw/tasks/fit.py`: `fit_to_model`, `restore_to_source` | fit_to_model's record is the whole of what restore_to_source needs: the scale is inferred from the output's size and must be uniform, and restore returns the source's size times that scale and at most its frame count. | `tests/test_fit_to_model.py` |
77-
| Shared image ops | `dw/tasks/image_ops.py`: `LUMA_WEIGHTS`, `luma`, `split_alpha`, `join_alpha`, `per_frame` | `grade`, `finish` and `lut` take luma (float64 Rec. 709 weights, cast to the input's dtype) and alpha handling from `image_ops`, never from each other, so an RGBA input comes back RGBA with its alpha untouched; `lut.LUMA_WEIGHTS` is a re-export. An image command runs over a video or frame list only through `per_frame`. | `tests/test_image_ops.py::TestLuma::test_lut_re_exports_the_one_float64_weights`, `tests/test_image_ops.py::TestAlpha::test_rgba_round_trip_keeps_the_alpha_bytes`, `tests/test_image_ops.py::TestOutputsAreByteIdentical::test_every_command_matches_the_pre_refactor_bytes`, `tests/test_image_ops.py::TestPerFrame::test_a_video_comes_back_as_a_video_with_the_same_frames_and_audio` |
77+
| Shared image ops | `dw/tasks/image_ops.py`: `LUMA_WEIGHTS`, `luma`, `split_alpha`, `join_alpha`, `per_frame` | `grade`, `finish` and `lut` take luma (float64 Rec. 709 weights, cast to the input's dtype) and alpha handling from `image_ops`, never from each other, so an RGBA input comes back RGBA with its alpha untouched; `lut.LUMA_WEIGHTS` is a re-export. An image command runs over a video or frame list only through `per_frame`. A media file is read by `task._load_media` with `fetch_image(keep_alpha=True)`, so a file with transparency arrives RGBA; a pipeline input still flattens to RGB. | `tests/test_image_ops.py::TestLuma::test_lut_re_exports_the_one_float64_weights`, `tests/test_image_ops.py::TestAlpha::test_rgba_round_trip_keeps_the_alpha_bytes`, `tests/test_image_ops.py::TestOutputsAreByteIdentical::test_every_command_matches_the_pre_refactor_bytes`, `tests/test_image_ops.py::TestPerFrame::test_a_video_comes_back_as_a_video_with_the_same_frames_and_audio`, `tests/test_finish.py::TestRgbaFileKeepsAlpha::test_a_file_path_comes_back_rgba_with_its_alpha` |
7878
| Audio level of a deliverable | `dw/audio_qc.py`: `warn_without_headroom`, `warn_if_written_above_full_scale` | A muxed video keeps the pre-encode headroom prediction and falls back to it only when the post-write probe cannot measure the file. A result a later step resets the level of (`consumed_by_normalizer`, `dw/step_cache.py`) draws neither the post-write warning nor the fallback. | `tests/test_result.py::TestNoHeadroom::test_a_clean_video_mux_drops_the_stale_prediction`, `tests/test_result.py::TestNoHeadroom::test_an_unprobeable_video_mux_falls_back_to_the_prediction`, `tests/test_result.py::TestNoHeadroom::test_a_video_a_join_level_matches_draws_no_headroom_warning`, `tests/test_result.py::TestTheWrittenLevel::test_a_file_that_decodes_above_full_scale_warns` |
7979
| Template audio-level placement | `workflows/templates/minimax/music-video.json`, `workflows/templates/minimax/music.json` | `normalize_audio` (-3 dBFS) acts only on the track that goes into the final file, so the slices that condition the shots are untouched. `music-video` slices each shot `lead_frames` early and `trim_video` drops the run-up before the join. | — |
8080

‎dw/argument_media.py‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,14 +115,26 @@ def _fetch_audio_video(video_spec, base_dir=None):
115115
return video_spec
116116

117117

118-
def fetch_image(img_spec, base_dir=None):
118+
def _rgb_keeping_alpha(image):
119+
"""RGBA when the file carries transparency (an alpha band, or a palette
120+
or greyscale image's transparency key), plain RGB otherwise."""
121+
if "A" in image.getbands() or "transparency" in image.info:
122+
return image.convert("RGBA")
123+
return image.convert("RGB")
124+
125+
126+
def fetch_image(img_spec, base_dir=None, keep_alpha=False):
119127
"""
120128
Load image from file path or URL with security validation.
121129
122130
Args:
123131
img_spec: Image specification (file path, URL, dict with 'location' key, PIL Image, or list of any of these)
124132
base_dir: Directory relative file paths are resolved against - the
125133
workflow file's directory. Defaults to the process working directory
134+
keep_alpha: Load a file with transparency as RGBA instead of
135+
flattening it to RGB. A pipeline input wants RGB; the finishing
136+
commands (grade, sharpen, film_grain, apply_lut) put the alpha
137+
back after their op, so a cutout stays a cutout (#775)
126138
127139
Returns:
128140
Loaded PIL Image, list of PIL Images, or None if img_spec is None
@@ -137,7 +149,7 @@ def fetch_image(img_spec, base_dir=None):
137149
# Handle lists of images (recursively process each)
138150
if isinstance(img_spec, list):
139151
logger.debug(f"Loading list of {len(img_spec)} images")
140-
return [fetch_image(img, base_dir) for img in img_spec]
152+
return [fetch_image(img, base_dir, keep_alpha) for img in img_spec]
141153

142154
# If already a PIL Image, return as-is (allows multiple realize_args calls)
143155
if hasattr(img_spec, "mode") and hasattr(img_spec, "size"):
@@ -161,6 +173,7 @@ def fetch_image(img_spec, base_dir=None):
161173
return img_spec
162174

163175
logger.debug(f"Loading image from: {img_spec}")
176+
convert_method = _rgb_keeping_alpha if keep_alpha else None
164177

165178
try:
166179
# Check if it's a URL
@@ -171,7 +184,7 @@ def fetch_image(img_spec, base_dir=None):
171184
# without re-checking them; load_image still does the EXIF
172185
# transpose and RGB conversion on the decoded result
173186
response = safe_get(img_spec, "an image argument", timeout=60)
174-
return load_image(Image.open(io.BytesIO(response.content)))
187+
return load_image(Image.open(io.BytesIO(response.content)), convert_method)
175188
else:
176189
# Treat as file path, relative to the workflow file, and confined
177190
# to the directories this workflow may read (dw/locations.py)
@@ -185,7 +198,7 @@ def fetch_image(img_spec, base_dir=None):
185198
f"Image file extension not allowed: {ext} - a video file "
186199
"must go through video_frames first"
187200
)
188-
return load_image(validated_path)
201+
return load_image(validated_path, convert_method)
189202

190203
except SecurityError:
191204
raise

‎dw/tasks/task.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -872,7 +872,8 @@ def _load_media(media):
872872
a video extension is read with its audio, anything else as an image. A
873873
value that is not a string - a PIL Image, an AudioVideo, a frame list -
874874
is already loaded and comes back as itself. Shared by the finishing
875-
commands that take `media` (#603).
875+
commands that take `media` (#603). An image file with transparency loads
876+
as RGBA, so the command puts its alpha back after the op (#775).
876877
"""
877878
if not isinstance(media, str):
878879
return media
@@ -885,7 +886,7 @@ def _load_media(media):
885886
return load_audio_video(media)
886887
from ..argument_media import fetch_image
887888

888-
return fetch_image(media)
889+
return fetch_image(media, keep_alpha=True)
889890

890891

891892
@register_command(

‎tests/test_finish.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,3 +403,59 @@ def test_the_new_domains_are_described_with_their_range_text(self):
403403

404404
if __name__ == "__main__":
405405
pytest.main([__file__, "-v"])
406+
407+
408+
class TestRgbaFileKeepsAlpha:
409+
"""An RGBA file on disk - an asset:/output: path once resolved - loads with
410+
its alpha, so the finishing commands hand it back untouched (#775). The
411+
pipeline-input loader still flattens to RGB."""
412+
413+
@staticmethod
414+
def _cutout(tmp_path):
415+
rgb = numpy.full((16, 16, 3), 200, dtype=numpy.uint8)
416+
alpha = numpy.zeros((16, 16), dtype=numpy.uint8)
417+
alpha[4:12, 4:12] = 255
418+
image = Image.fromarray(rgb).convert("RGBA")
419+
image.putalpha(Image.fromarray(alpha))
420+
path = tmp_path / "cutout.png"
421+
image.save(path)
422+
return str(path), alpha
423+
424+
@pytest.mark.parametrize(
425+
"command, arguments",
426+
[
427+
("grade", {"exposure": 0.5, "saturation": 0.5}),
428+
("film_grain", {"amount": 0.3, "seed": 1}),
429+
("apply_lut", {"palette": ["#102030", "#e0c090"]}),
430+
("sharpen", {"amount": 1.0}),
431+
],
432+
)
433+
def test_a_file_path_comes_back_rgba_with_its_alpha(
434+
self, tmp_path, command, arguments
435+
):
436+
path, alpha = self._cutout(tmp_path)
437+
result = _run(command, path, **arguments)
438+
assert result.mode == "RGBA"
439+
assert numpy.array_equal(_array(result.getchannel("A")), alpha)
440+
441+
def test_an_opaque_file_still_loads_rgb(self, tmp_path):
442+
from dw.tasks.task import _load_media
443+
444+
path = tmp_path / "opaque.png"
445+
Image.new("RGB", (8, 8), "red").save(path)
446+
assert _load_media(str(path)).mode == "RGB"
447+
448+
def test_a_palette_file_with_transparency_loads_rgba(self, tmp_path):
449+
from dw.tasks.task import _load_media
450+
451+
path = tmp_path / "palette.png"
452+
image = Image.new("P", (8, 8), 0)
453+
image.info["transparency"] = 0
454+
image.save(path, transparency=0)
455+
assert _load_media(str(path)).mode == "RGBA"
456+
457+
def test_a_pipeline_input_still_flattens_to_rgb(self, tmp_path):
458+
from dw.argument_media import fetch_image
459+
460+
path, _ = self._cutout(tmp_path)
461+
assert fetch_image(path).mode == "RGB"

0 commit comments

Comments
 (0)

Back | FazBrowse Home | New Git URL