diff --git a/CLAUDE.md b/CLAUDE.md index 84dd5b35..1796bc16 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -27,7 +27,7 @@ python -m dw.test # Interactive REPL python -m dw.repl -# Run all tests (134+ tests) +# Run all tests (260+ tests) pytest -v # Run a single test file or test @@ -52,7 +52,7 @@ JSON workflow → schema validation → variable substitution → sequential ste 1. `workflow.py` loads JSON, validates against `workflow_schema.json`, substitutes `variable:name` references 2. `step.py` executes each step — generating argument combinations via `previous_results.py` (cartesian product of `previous_result:step_name` references) 3. Each step dispatches to one of: **Pipeline** (HuggingFace inference), **Task** (utility operation), or **Sub-Workflow** (recursive) -4. `result.py` saves outputs as `{output_dir}/{workflow_id}-{step_name}.{index}.{ext}` +4. `result.py` saves outputs as `{output_dir}/{workflow_id}-{step_name}.{index}.{ext}` — supports image, video, audio, text, and JSON content types. Optional `embed_metadata` stores generation parameters in PNG info chunks or JPEG/WebP EXIF. ### Key Modules @@ -62,7 +62,12 @@ JSON workflow → schema validation → variable substitution → sequential ste | `dw/step.py` | Step executor: generates iterations, dispatches to pipeline/task/workflow | | `dw/pipeline_processors/pipeline.py` | Pipeline loading, components, quantization, LoRA, schedulers, offloading | | `dw/pipeline_processors/config_objects.py` | Quantization and group offload config creation | -| `dw/tasks/task.py` | Task dispatcher (image processing, QR codes, gathering, video) | +| `dw/tasks/task.py` | Task dispatcher (image processing, QR codes, gathering, video, segmentation, captioning, frame interpolation) | +| `dw/tasks/segment.py` | GroundingDINO + SAM2 text-prompted object segmentation | +| `dw/tasks/image_to_text.py` | Image captioning via transformers image-to-text pipeline (BLIP, BLIP-2, etc.) | +| `dw/tasks/text_generation.py` | Text generation / prompt expansion via transformers text-generation pipeline | +| `dw/tasks/interpolate_frames.py` | RIFE frame interpolation (2x/4x/8x) with vendored IFNet v4.6 | +| `dw/tasks/rife_model.py` | Vendored RIFE IFNet v4.6 architecture (MIT License, Megvii Inc.) | | `dw/previous_results.py` | Cross-step data flow via cartesian products | | `dw/arguments.py` | Argument processing, resource loading, dynamic type conversion | | `dw/type_helpers.py` | Dynamic type loading: `"FluxPipeline"` → class, `"torch.bfloat16"` → dtype | diff --git a/README.md b/README.md index 4e31adb6..71139ecf 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,8 @@ A declarative workflow engine for the [Hugging Face Diffusers library](https://g - **Prompt weighting** — A1111-style `(word:1.5)` syntax with long prompt support - **LoRA and IP-Adapter** support - **Composable workflows** from multiple JSON files with `builtin:` references -- **Utility tasks** — background removal, upscaling, cropping, QR codes, LLM prompt augmentation +- **Utility tasks** — upscaling, face restoration, segmentation, captioning, frame interpolation, QR codes, and more +- **Metadata embedding** — store generation parameters in PNG/JPEG/WebP for reproducibility - **Interactive REPL** with persistent GPU model caching (2-4x faster iteration) - **Cross-platform** — CUDA, MPS (Apple Silicon), and CPU diff --git a/TODO.md b/TODO.md new file mode 100644 index 00000000..a3cb4bd1 --- /dev/null +++ b/TODO.md @@ -0,0 +1,47 @@ +# diffusers-workflow TODO + +Utilities and enhancements from the stable diffusion ecosystem. + +## Post-Processing & Enhancement + +- [x] **CCSR / StableSR** — Diffusion-based upscalers with better detail than ESRGAN-family, especially for faces and textures. New task type alongside existing Spandrel upscaler. *(Implemented as `diffusion_upscale` task wrapping StableDiffusionUpscalePipeline x4/x2)* +- [ ] **Real-ESRGAN Video** — Frame-consistent video upscaling with temporal smoothing. Current upscaler is image-only. +- [x] **RIFE / FILM frame interpolation** — Generate intermediate frames for smoother video output. Video post-processing task. + +## Image Preprocessing & Conditioning + +- [x] **Marigold depth / DSINE normals** — Newer, more accurate depth/normal estimators than MiDaS/DPT. Better ControlNet conditioning maps. +- [x] **GroundingDINO + SAM2** — Text-prompted object detection to segmentation. "Segment the dog" as a task, producing masks for inpainting workflows. +- [x] **Florence-2** — Microsoft vision-language model for captioning, detection, segmentation. Powers an `auto_caption` task for img2img or IP-Adapter workflows. +- [ ] **PuLID / InstantID** — Identity-preserving face conditioning (better than IP-Adapter for faces). Works with Flux and SDXL. + +## Video-Specific + +- [ ] **FramePack** — Context-aware video generation with efficient memory usage for long video generation. +- [ ] **PySceneDetect keyframe extraction** — Smarter than `get_frame` for selecting keyframes from input video based on scene detection. +- [x] **RIFE/IFRNet optical flow interpolation** — Optical flow-based frame interpolation as a post-processing step. + +## Workflow Utilities + +- [x] **Prompt expansion via local LLM** — Task that takes a short prompt and expands it using a small language model (Llama 3.2 1B, Qwen2.5, etc.). +- [ ] **Image comparison / SSIM / LPIPS scoring** — Task that scores similarity between images for iterative refinement workflows. +- [ ] **Color palette extraction / transfer** — Extract dominant colors from a reference image or apply color grading from one image to another. +- [ ] **Tiled generation / outpainting helper** — Automate the border+mask+generation loop for progressive outpainting. + +## Model Management + +- [ ] **Automatic VRAM estimation** — Given a workflow JSON, estimate peak VRAM before running. Helps pick the right quantization/offload settings. +- [ ] **Model predownload/warmup** — Dry-run mode that downloads all models without executing. Useful for deployment. + +## Quick Wins + +- [x] **Image metadata embedding** — Store generation params in PNG info chunks for reproducibility. +- [x] **EXIF stripping** on input images — Privacy-safe preprocessing. +- [ ] **Image hashing** (perceptual hash) — Dedup detection across workflow runs. +- [x] **Watermark embedding/detection** — Responsible AI compliance. +- [x] **Aspect ratio bucketing** — Auto-resize inputs to model-native aspect ratios. + +## Architectural Enhancements + +- [ ] **Conditional branching** — e.g., "if image has faces, run face restore; otherwise skip." Enables more sophisticated pipelines. +- [ ] **Parallel step execution** — Steps with no data dependencies run concurrently. Matters for multi-GPU setups. diff --git a/docs/TASKS.md b/docs/TASKS.md index cad77713..3e3160e2 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -72,11 +72,87 @@ All accept an `image` argument with processing parameters: | `resize_center_crop` | Resize with center crop | `width`, `height` | | `resize_resample` | Resample to nearest 64px multiple | | | `resize_rescale` | Resize to exact dimensions | `width`, `height` | +| `resize_bucket` | Snap to closest model-native aspect ratio | `resolution`, `ratios`, `alignment` | | `crop_square` | Center crop to square | | | `add_border_and_mask` | Add border with alpha mask | | | `add_border_and_mask_with_size` | Border with specific dimensions | `width`, `height` | +| `strip_exif` | Remove all EXIF/metadata from image | | +| `add_watermark` | Add visible text watermark | `text`, `position`, `opacity`, `font_size`, `color`, `margin` | | `get_image_size` | Return `{width, height}` dict | | +### EXIF Stripping + +Remove all EXIF metadata, GPS coordinates, camera info, and timestamps from images for privacy-safe preprocessing: + +```json +{ + "task": { + "command": "strip_exif", + "arguments": { + "image": "previous_result:input_image" + } + }, + "result": { "content_type": "image/png" } +} +``` + +Returns a clean copy with pixel data only — no embedded metadata. Useful as a first step when processing user-uploaded images. + +### Watermark Embedding + +Add a visible text watermark to images for responsible AI compliance: + +```json +{ + "task": { + "command": "add_watermark", + "arguments": { + "image": "previous_result:generate", + "text": "AI Generated", + "position": "bottom-right", + "opacity": 128 + } + }, + "result": { "content_type": "image/png" } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `text` | No | Watermark text (default: "AI Generated") | +| `position` | No | "bottom-right", "bottom-left", "top-right", "top-left", or "center" (default: "bottom-right") | +| `opacity` | No | Text opacity 0-255 (default: 128) | +| `font_size` | No | Font size in pixels, 0 = auto-scale ~3% of image height (default: 0) | +| `color` | No | RGB array for text color (default: white) | +| `margin` | No | Pixel margin from edges (default: 10) | + +### Aspect Ratio Bucketing + +The `resize_bucket` command snaps an image to the closest model-native aspect ratio, then resizes with 64-pixel alignment. This avoids distortion and ensures the model generates at a resolution it was trained on. + +```json +{ + "task": { + "command": "resize_bucket", + "arguments": { + "image": "previous_result:input_image", + "resolution": 1024 + } + }, + "result": { "content_type": "image/png" } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `resolution` | No | Target short-side size in pixels (default: 1024) | +| `ratios` | No | Custom list of `[w, h]` ratio pairs (default: standard SDXL/Flux ratios) | +| `alignment` | No | Round dimensions to this multiple (default: 64) | + +**Default ratios:** 1:1, 4:3, 3:4, 3:2, 2:3, 16:9, 9:16, 21:9, 9:21 + +For example, a 1600x900 photo (16:9) at resolution 1024 becomes 1792x1024. A 800x600 photo (4:3) becomes 1344x1024. + ## Video Processing | Command | Description | Extra Arguments | @@ -142,6 +218,43 @@ Large images are automatically tiled to avoid GPU memory issues. Models can be l **Example:** [SpandrelUpscale.json](../examples/SpandrelUpscale.json) — Generate at 512px, then 4x upscale to 2048px. +## Diffusion Upscaling + +Upscale images using Stable Diffusion upscale pipelines. Text-guided upscaling with better detail recovery than traditional super-resolution, especially for faces and textures. + +Two modes are available: +- **x4** (default): `StableDiffusionUpscalePipeline` — 4x upscale via `stabilityai/stable-diffusion-x4-upscaler` +- **x2**: `StableDiffusionLatentUpscalePipeline` — 2x upscale via `stabilityai/sd-x2-latent-upscaler` + +```json +{ + "task": { + "command": "diffusion_upscale", + "arguments": { + "image": "previous_result:generate", + "prompt": "high quality, detailed", + "negative_prompt": "blurry, low quality, artifacts", + "mode": "x4" + } + } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `image` | Yes | PIL Image or `previous_result:` reference | +| `prompt` | No | Text guidance for upscaling (default: "") | +| `negative_prompt` | No | Negative text guidance (default: none) | +| `mode` | No | `"x4"` or `"x2"` (default: `"x4"`) | +| `model_name` | No | Override the default model for the selected mode | +| `num_inference_steps` | No | Denoising steps (default: 25) | +| `guidance_scale` | No | Classifier-free guidance scale (default: 9.0) | +| `noise_level` | No | Noise level for x4 mode (default: 20, ignored for x2) | + +**Examples:** +- [DiffusionUpscaleX4.json](../examples/DiffusionUpscaleX4.json) — Generate at 512px, then 4x diffusion upscale to 2048px. +- [DiffusionUpscaleX2.json](../examples/DiffusionUpscaleX2.json) — Generate at 512px, then 2x latent upscale to 1024px. + ## Face Restoration Restore and enhance faces in images using spandrel-compatible face restoration models (GFPGAN, CodeFormer, RestoreFormer). Uses facexlib for face detection and alignment, then runs each detected face through the restoration model. @@ -220,6 +333,167 @@ You can chain upscaling and face restoration. Generate first, upscale the backgr This gives the best results: the super-resolution model handles background detail while the face model handles facial features, composited together at the upscaled resolution. +## Object Segmentation + +Detect and segment objects using text prompts via GroundingDINO + SAM2. Returns a binary mask image suitable for inpainting workflows. + +```json +{ + "task": { + "command": "segment", + "arguments": { + "image": "previous_result:input_image", + "prompt": "dog" + } + }, + "result": { "content_type": "image/png" } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `image` | Yes | PIL Image or `previous_result:` reference | +| `prompt` | Yes | Text description of object(s) to detect (e.g., "dog", "red car") | +| `model_name` | No | GroundingDINO model ID (default: `IDEA-Research/grounding-dino-base`) | +| `sam_model_name` | No | SAM2 model ID (default: `facebook/sam2-hiera-large`) | +| `threshold` | No | Detection confidence threshold (default: 0.3) | +| `invert` | No | Invert the output mask (default: false) | + +Returns a grayscale PIL Image (mode "L") — white (255) for detected objects, black (0) for background. Use with inpainting pipelines like FluxFillPipeline. + +**Examples:** + +- [Segment.json](../examples/Segment.json) — Segment an object from an image +- [SegmentAndInpaint.json](../examples/SegmentAndInpaint.json) — Segment, then inpaint the masked region + +## Image Captioning + +Generate text captions from images using HuggingFace image-to-text models (BLIP, BLIP-2, ViT-GPT2, GIT, etc.). + +```json +{ + "task": { + "command": "image_to_text", + "arguments": { + "image": "previous_result:input_image" + } + }, + "result": { "content_type": "text/plain" } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `image` | Yes | PIL Image or `previous_result:` reference | +| `model_name` | No | HuggingFace model ID (default: `Salesforce/blip-image-captioning-base`) | +| `prompt` | No | Text prompt for conditional captioning (supported by BLIP-2, etc.) | +| `max_new_tokens` | No | Maximum tokens to generate (default: 50) | + +Returns a caption string. Save as `text/plain` for `.txt` output, or pass to a downstream step via `previous_result:` as a prompt for image generation. + +For Florence-2's advanced task-token captioning (detailed captions, object detection, OCR), use the built-in `describe_image` workflow instead: + +```json +{ + "name": "caption", + "workflow": { + "path": "builtin:describe_image.json", + "arguments": { "image": "previous_result:input_image" } + }, + "result": { "content_type": "text/plain" } +} +``` + +**Examples:** + +- [ImageToText.json](../examples/ImageToText.json) — Basic BLIP captioning, saves as `.txt` +- [ImageToTextBlip2.json](../examples/ImageToTextBlip2.json) — BLIP-2 with conditional prompt +- [CaptionToImage.json](../examples/CaptionToImage.json) — Caption an image, then regenerate with Flux + +## Text Generation / Prompt Expansion + +Generate or expand text using a local language model. Useful for expanding short prompts into detailed image generation prompts, rewriting text, or other text-to-text tasks. + +```json +{ + "task": { + "command": "text_generation", + "arguments": { + "prompt": "a cat on a windowsill", + "system_prompt": "You are a helpful AI assistant that creates detailed prompts for text to image generative AI. When supplied input generate only the prompt, no other text." + } + }, + "result": { "content_type": "text/plain" } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `prompt` | Yes | The user message or short prompt to expand/transform | +| `system_prompt` | No | System instruction for the model (e.g., "expand this into a detailed image prompt") | +| `model_name` | No | HuggingFace model ID (default: `Qwen/Qwen2.5-1.5B-Instruct`) | +| `max_new_tokens` | No | Maximum tokens to generate (default: 500) | + +Returns a text string. Save as `text/plain` for `.txt` output, or pass to a downstream step via `previous_result:` as a prompt for image generation. + +Any HuggingFace chat model works — Qwen2.5, Llama 3.2, Phi-3.5, etc. The default (Qwen2.5-1.5B-Instruct) is small enough to run alongside diffusion models. + +There is also a built-in `augment_prompt` workflow (`builtin:augment_prompt.json`) that does the same thing using a 3-step pipeline approach with Phi-3.5-mini. The `text_generation` task is the simpler single-step alternative. + +**Examples:** + +- [ExpandPrompt.json](../examples/ExpandPrompt.json) — Expand a short prompt and save as `.txt` +- [ExpandAndGenerate.json](../examples/ExpandAndGenerate.json) — Expand prompt, then generate with Flux + +## Frame Interpolation + +Increase video frame rate using RIFE (Real-Time Intermediate Flow Estimation). Takes a list of video frames and inserts intermediate frames between each pair. + +```json +{ + "task": { + "command": "interpolate_frames", + "arguments": { + "video": "previous_result:generate_video", + "multiplier": 2 + } + }, + "result": { "content_type": "video/mp4", "fps": 60 } +} +``` + +| Argument | Required | Description | +| -------- | -------- | ----------- | +| `video` | Yes | List of PIL Images (video frames) or `previous_result:` reference | +| `multiplier` | No | Frame count multiplier: 2, 4, or 8 (default: 2) | +| `model_name` | No | HuggingFace repo with RIFE weights (default: `styler00dollar/RIFE-v4.6`) | + +Uses vendored IFNet v4.6 architecture. Weights are downloaded from HuggingFace Hub on first use. + +**Example:** [InterpolateFrames.json](../examples/InterpolateFrames.json) — Generate video with Mochi, then 2x interpolate from 30fps to 60fps. + +## Metadata Embedding + +Embed generation parameters in saved images. Enable by setting `embed_metadata: true` in a step's result configuration: + +```json +{ + "result": { + "content_type": "image/png", + "embed_metadata": true + } +} +``` + +| Format | Storage | Notes | +| ------ | ------- | ----- | +| PNG | Text chunk (`parameters` key) | Always available | +| JPEG/WebP | EXIF UserComment | Requires `pip install piexif` | + +Metadata includes step name, model name, and generation arguments (prompt, steps, guidance scale, etc.) as JSON. + +**Example:** [MetadataEmbed.json](../examples/MetadataEmbed.json) — Generate with Flux and embed parameters in PNG. + ## QR Code Generation ```json @@ -285,3 +559,12 @@ Canny edge detection followed by ControlNet generation: - [upscale.json](../examples/upscale.json) — Gather, resize, and diffusion upscale - [SpandrelUpscale.json](../examples/SpandrelUpscale.json) — Generate + spandrel 4x upscale - [FaceRestore.json](../examples/FaceRestore.json) — Generate portrait + GFPGAN face restoration +- [Segment.json](../examples/Segment.json) — Text-prompted object segmentation +- [SegmentAndInpaint.json](../examples/SegmentAndInpaint.json) — Segment + inpaint +- [ImageToText.json](../examples/ImageToText.json) — BLIP image captioning +- [ImageToTextBlip2.json](../examples/ImageToTextBlip2.json) — BLIP-2 conditional captioning +- [CaptionToImage.json](../examples/CaptionToImage.json) — Caption then regenerate +- [InterpolateFrames.json](../examples/InterpolateFrames.json) — RIFE frame interpolation +- [MetadataEmbed.json](../examples/MetadataEmbed.json) — Embed generation parameters in PNG +- [ExpandPrompt.json](../examples/ExpandPrompt.json) — LLM prompt expansion +- [ExpandAndGenerate.json](../examples/ExpandAndGenerate.json) — Expand prompt + generate image diff --git a/docs/superpowers/plans/2026-03-26-ecosystem-utilities.md b/docs/superpowers/plans/2026-03-26-ecosystem-utilities.md new file mode 100644 index 00000000..d2da6433 --- /dev/null +++ b/docs/superpowers/plans/2026-03-26-ecosystem-utilities.md @@ -0,0 +1,1377 @@ +# Ecosystem Utilities Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add Marigold depth/normals example workflows, a GroundingDINO+SAM2 segmentation task, a RIFE frame interpolation task, and opt-in image metadata embedding. + +**Architecture:** Four independent features following existing patterns. Two new task commands (`segment`, `interpolate_frames`) registered via `@register_command` in `task.py` with implementations in separate files. Metadata embedding modifies the existing `Result.save_artifact` path. Marigold is example-only (no code). + +**Tech Stack:** Python, PIL, PyTorch, transformers (GroundingDINO/SAM2), diffusers (Marigold), piexif (optional, JPEG metadata) + +--- + +## File Structure + +| File | Action | Responsibility | +|------|--------|---------------| +| `examples/MarigoldDepth.json` | Create | Example: gather image then MarigoldDepthPipeline | +| `examples/MarigoldNormals.json` | Create | Example: gather image then MarigoldNormalsPipeline | +| `dw/tasks/segment.py` | Create | GroundingDINO+SAM2 segmentation implementation | +| `dw/tasks/task.py:1-14,84-99` | Modify | Import segment + interpolate, register commands | +| `tests/test_segment.py` | Create | Unit tests for segment task | +| `examples/Segment.json` | Create | Example: gather image then segment then save mask | +| `examples/SegmentAndInpaint.json` | Create | Example: gather then segment then FluxFill inpaint | +| `dw/tasks/interpolate_frames.py` | Create | RIFE frame interpolation implementation | +| `tests/test_interpolate_frames.py` | Create | Unit tests for interpolation task | +| `examples/InterpolateFrames.json` | Create | Example: Mochi video then interpolate then save | +| `dw/result.py:17-35,159-238` | Modify | Add metadata support to Result class and save_artifact | +| `dw/step.py:26-96` | Modify | Collect and pass metadata to Result | +| `dw/workflow_schema.json:650-689` | Modify | Add embed_metadata property to result definition | +| `tests/test_result.py` | Modify | Add metadata embedding tests | +| `examples/MetadataEmbed.json` | Create | Example: Flux generate with embed_metadata: true | + +--- + +### Task 1: Marigold Depth Example Workflow + +**Files:** +- Create: `examples/MarigoldDepth.json` + +- [ ] **Step 1: Create MarigoldDepth.json** + +```json +{ + "id": "MarigoldDepth", + "variables": { + "image_url": "https://marigoldmonodepth.github.io/images/einstein.jpg" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "depth", + "pipeline": { + "configuration": { + "component_type": "MarigoldDepthPipeline" + }, + "from_pretrained_arguments": { + "model_name": "prs-eth/marigold-depth-lcm-v1-0", + "torch_dtype": "torch.float16", + "variant": "fp16" + }, + "arguments": { + "image": "previous_result:load_image" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} +``` + +- [ ] **Step 2: Validate the workflow JSON against the schema** + +Run: `python -m dw.validate examples/MarigoldDepth.json` +Expected: Validation passes with no errors. + +- [ ] **Step 3: Commit** + +```bash +git add examples/MarigoldDepth.json +git commit -m "feat: add MarigoldDepth example workflow" +``` + +--- + +### Task 2: Marigold Normals Example Workflow + +**Files:** +- Create: `examples/MarigoldNormals.json` + +- [ ] **Step 1: Create MarigoldNormals.json** + +```json +{ + "id": "MarigoldNormals", + "variables": { + "image_url": "https://marigoldmonodepth.github.io/images/einstein.jpg" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "normals", + "pipeline": { + "configuration": { + "component_type": "MarigoldNormalsPipeline" + }, + "from_pretrained_arguments": { + "model_name": "prs-eth/marigold-normals-lcm-v1-0", + "torch_dtype": "torch.float16", + "variant": "fp16" + }, + "arguments": { + "image": "previous_result:load_image" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} +``` + +- [ ] **Step 2: Validate the workflow JSON against the schema** + +Run: `python -m dw.validate examples/MarigoldNormals.json` +Expected: Validation passes with no errors. + +- [ ] **Step 3: Commit** + +```bash +git add examples/MarigoldNormals.json +git commit -m "feat: add MarigoldNormals example workflow" +``` + +--- + +### Task 3: Segment Task — Tests + +**Files:** +- Create: `tests/test_segment.py` + +- [ ] **Step 1: Write the unit tests for segment_image** + +These tests mock the transformers models to avoid downloading multi-GB weights in CI. + +```python +import pytest +from unittest.mock import patch, MagicMock +import torch +import numpy as np +from PIL import Image + + +def _make_test_image(width=640, height=480): + """Create a simple test image.""" + return Image.new("RGB", (width, height), color=(128, 64, 32)) + + +class TestSegmentImage: + """Test segment_image function with mocked models.""" + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_returns_pil_image_mode_l( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + """segment_image should return a grayscale PIL Image.""" + from dw.tasks.segment import segment_image + + # Mock GroundingDINO detection + mock_processor_instance = MagicMock() + mock_auto_proc.from_pretrained.return_value = mock_processor_instance + mock_processor_instance.return_value = {"input_ids": torch.zeros(1, 10)} + mock_processor_instance.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.tensor([[100.0, 100.0, 300.0, 300.0]]), + "scores": torch.tensor([0.9]), + "labels": ["dog"], + } + ] + + mock_model_instance = MagicMock() + mock_auto_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.return_value = MagicMock() + + # Mock SAM2 + mock_sam_proc_instance = MagicMock() + mock_sam_proc.from_pretrained.return_value = mock_sam_proc_instance + mock_sam_proc_instance.return_value = { + "pixel_values": torch.zeros(1, 3, 256, 256), + "input_boxes": torch.tensor([[[100.0, 100.0, 300.0, 300.0]]]), + } + + mock_sam_model_instance = MagicMock() + mock_sam_model.from_pretrained.return_value = mock_sam_model_instance + # SAM2 output: pred_masks shape [1, 1, 3, H, W] + mask_tensor = torch.zeros(1, 1, 3, 480, 640) + mask_tensor[0, 0, 0, 100:300, 100:300] = 1.0 + mock_sam_model_instance.return_value = MagicMock(pred_masks=mask_tensor) + mock_sam_proc_instance.post_process_masks.return_value = [ + torch.zeros(1, 1, 480, 640) + ] + # Set the mask area to 1 + mock_sam_proc_instance.post_process_masks.return_value[0][0, 0, 100:300, 100:300] = 1.0 + + image = _make_test_image() + result = segment_image(image, "dog") + + assert isinstance(result, Image.Image) + assert result.mode == "L" + assert result.size == (640, 480) + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_no_detections_returns_black_mask( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + """When nothing is detected, return an all-black mask.""" + from dw.tasks.segment import segment_image + + mock_processor_instance = MagicMock() + mock_auto_proc.from_pretrained.return_value = mock_processor_instance + mock_processor_instance.return_value = {"input_ids": torch.zeros(1, 10)} + mock_processor_instance.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.zeros(0, 4), + "scores": torch.zeros(0), + "labels": [], + } + ] + + mock_model_instance = MagicMock() + mock_auto_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.return_value = MagicMock() + + image = _make_test_image() + result = segment_image(image, "nonexistent_object") + + assert isinstance(result, Image.Image) + assert result.mode == "L" + # All black = no detection + arr = np.array(result) + assert arr.max() == 0 + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_invert_flag( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + """When invert=True, mask should be inverted.""" + from dw.tasks.segment import segment_image + + mock_processor_instance = MagicMock() + mock_auto_proc.from_pretrained.return_value = mock_processor_instance + mock_processor_instance.return_value = {"input_ids": torch.zeros(1, 10)} + mock_processor_instance.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.zeros(0, 4), + "scores": torch.zeros(0), + "labels": [], + } + ] + + mock_model_instance = MagicMock() + mock_auto_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.return_value = MagicMock() + + image = _make_test_image() + result = segment_image(image, "nonexistent_object", invert=True) + + assert isinstance(result, Image.Image) + # Inverted black mask = all white + arr = np.array(result) + assert arr.min() == 255 + + +class TestSegmentTaskRegistration: + """Test that segment is properly registered as a task command.""" + + def test_segment_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + assert "segment" in _COMMAND_REGISTRY +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/test_segment.py -v` +Expected: FAIL — `dw.tasks.segment` module does not exist yet. + +- [ ] **Step 3: Commit** + +```bash +git add tests/test_segment.py +git commit -m "test: add segment task tests" +``` + +--- + +### Task 4: Segment Task — Implementation + +**Files:** +- Create: `dw/tasks/segment.py` +- Modify: `dw/tasks/task.py:1-14` (imports), `dw/tasks/task.py:99` (after restore_faces handler) + +- [ ] **Step 1: Create dw/tasks/segment.py** + +```python +""" +Image segmentation via GroundingDINO + SAM2. + +Takes an image and a text prompt, detects objects matching the prompt, +and returns a binary mask image (white = detected object). +""" + +import logging +import torch +import numpy as np +from PIL import Image +from transformers import ( + AutoProcessor, + AutoModelForZeroShotObjectDetection, + Sam2Model, + Sam2Processor, +) + +logger = logging.getLogger("dw") + +# Default model IDs +_DEFAULT_DINO_MODEL = "IDEA-Research/grounding-dino-base" +_DEFAULT_SAM_MODEL = "facebook/sam2-hiera-large" + + +def segment_image(image, prompt, device="cpu", **kwargs): + """Segment objects matching a text prompt, returning a binary mask. + + Args: + image: PIL Image to segment + prompt: Text description of object(s) to detect (e.g., "dog") + device: Target device ("cuda", "mps", "cpu") + **kwargs: + model_name: GroundingDINO model ID (default: IDEA-Research/grounding-dino-base) + sam_model_name: SAM2 model ID (default: facebook/sam2-hiera-large) + threshold: Detection confidence threshold (default: 0.3) + invert: Invert the output mask (default: False) + + Returns: + PIL Image in mode "L" — white (255) for detected objects, black (0) for background. + """ + model_name = kwargs.get("model_name", _DEFAULT_DINO_MODEL) + sam_model_name = kwargs.get("sam_model_name", _DEFAULT_SAM_MODEL) + threshold = kwargs.get("threshold", 0.3) + invert = kwargs.get("invert", False) + + width, height = image.size + + # --- GroundingDINO: detect bounding boxes --- + logger.info(f"Loading GroundingDINO from {model_name}") + dino_processor = AutoProcessor.from_pretrained(model_name) + dino_model = AutoModelForZeroShotObjectDetection.from_pretrained(model_name).to( + device + ) + + inputs = dino_processor(images=image, text=prompt, return_tensors="pt").to(device) + + with torch.inference_mode(): + outputs = dino_model(**inputs) + + results = dino_processor.post_process_grounded_object_detection( + outputs, + inputs["input_ids"], + threshold=threshold, + target_sizes=[(height, width)], + ) + + boxes = results[0]["boxes"] # shape: [N, 4] + scores = results[0]["scores"] + labels = results[0]["labels"] + + logger.info(f"Detected {len(boxes)} objects: {labels} (scores: {scores.tolist()})") + + # If no detections, return blank mask + if len(boxes) == 0: + mask_image = Image.new("L", (width, height), 0) + if invert: + mask_image = Image.eval(mask_image, lambda x: 255 - x) + return mask_image + + # --- SAM2: generate masks from boxes --- + logger.info(f"Loading SAM2 from {sam_model_name}") + sam_processor = Sam2Processor.from_pretrained(sam_model_name) + sam_model = Sam2Model.from_pretrained(sam_model_name).to(device) + + # Format boxes for SAM2: [[[x1, y1, x2, y2], ...]] + input_boxes = [boxes.cpu().tolist()] + + sam_inputs = sam_processor( + images=image, + input_boxes=input_boxes, + return_tensors="pt", + ).to(device) + + with torch.inference_mode(): + sam_outputs = sam_model(**sam_inputs) + + masks = sam_processor.post_process_masks( + sam_outputs.pred_masks, + sam_inputs["original_sizes"], + sam_inputs["reshaped_input_sizes"], + ) + + # Combine all masks via union (logical OR) + # masks[0] shape: [N, 1, H, W] — take best mask per detection + combined = masks[0][:, 0].sum(dim=0).clamp(0, 1) # [H, W] + mask_array = (combined.cpu().numpy() * 255).astype(np.uint8) + + mask_image = Image.fromarray(mask_array, mode="L") + + if invert: + mask_image = Image.eval(mask_image, lambda x: 255 - x) + + logger.info(f"Generated segmentation mask {width}x{height}") + return mask_image +``` + +- [ ] **Step 2: Register the segment command in task.py** + +Add import at the top of `dw/tasks/task.py`, after the `restore_faces` import (line 13): + +```python +from .segment import segment_image +``` + +Add handler after the `_handle_restore_faces` function (after line 99): + +```python +@register_command("segment") +def _handle_segment(task, arguments, previous_pipelines): + """Segment objects in an image using text prompt""" + logger.debug("Segmenting image") + image = arguments.pop("image") + prompt = arguments.pop("prompt") + return segment_image(image, prompt, device=task.device, **arguments) +``` + +- [ ] **Step 3: Run tests to verify they pass** + +Run: `pytest tests/test_segment.py -v` +Expected: All tests PASS. + +- [ ] **Step 4: Run full test suite to check for regressions** + +Run: `pytest -v` +Expected: All existing tests continue to pass. + +- [ ] **Step 5: Commit** + +```bash +git add dw/tasks/segment.py dw/tasks/task.py +git commit -m "feat: add segment task (GroundingDINO + SAM2)" +``` + +--- + +### Task 5: Segment Example Workflows + +**Files:** +- Create: `examples/Segment.json` +- Create: `examples/SegmentAndInpaint.json` + +- [ ] **Step 1: Create examples/Segment.json** + +```json +{ + "id": "Segment", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/grounding_dino_example_input.png", + "prompt": "cat" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "segment", + "task": { + "command": "segment", + "arguments": { + "image": "previous_result:load_image", + "prompt": "variable:prompt" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} +``` + +- [ ] **Step 2: Create examples/SegmentAndInpaint.json** + +```json +{ + "id": "SegmentAndInpaint", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/grounding_dino_example_input.png", + "segment_prompt": "cat", + "inpaint_prompt": "a golden retriever puppy sitting on the grass" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "segment", + "task": { + "command": "segment", + "arguments": { + "image": "previous_result:load_image", + "prompt": "variable:segment_prompt" + } + }, + "result": { + "content_type": "image/png", + "save": true + } + }, + { + "name": "inpaint", + "pipeline": { + "configuration": { + "component_type": "FluxFillPipeline", + "offload": "sequential" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-Fill-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "image": "previous_result:load_image", + "mask_image": "previous_result:segment", + "prompt": "variable:inpaint_prompt", + "height": 1024, + "width": 1024, + "guidance_scale": 30, + "num_inference_steps": 50, + "max_sequence_length": 512 + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} +``` + +- [ ] **Step 3: Validate both workflows against the schema** + +Run: `python -m dw.validate examples/Segment.json && python -m dw.validate examples/SegmentAndInpaint.json` +Expected: Both pass validation. + +- [ ] **Step 4: Commit** + +```bash +git add examples/Segment.json examples/SegmentAndInpaint.json +git commit -m "feat: add Segment and SegmentAndInpaint example workflows" +``` + +--- + +### Task 6: Frame Interpolation Task — Tests + +**Files:** +- Create: `tests/test_interpolate_frames.py` + +- [ ] **Step 1: Write the unit tests for interpolate_frames** + +These tests mock the RIFE model to avoid downloading weights. They test the frame list logic. + +```python +import pytest +from unittest.mock import patch, MagicMock +from PIL import Image +import numpy as np + + +def _make_test_frames(count=4, width=64, height=64): + """Create a list of test frames with different colors.""" + frames = [] + for i in range(count): + shade = int(255 * i / max(count - 1, 1)) + frames.append(Image.new("RGB", (width, height), color=(shade, shade, shade))) + return frames + + +class TestInterpolateFrames: + """Test interpolate_frames function.""" + + @patch("dw.tasks.interpolate_frames._load_rife_model") + def test_2x_doubles_frame_count(self, mock_load): + """2x multiplier should produce 2N-1 frames from N input frames.""" + from dw.tasks.interpolate_frames import interpolate_frames + + # Mock model: return average of two input frames + mock_model = MagicMock() + + def fake_inference(img1, img2): + arr1 = np.array(img1).astype(np.float32) + arr2 = np.array(img2).astype(np.float32) + mid = ((arr1 + arr2) / 2).astype(np.uint8) + return Image.fromarray(mid) + + mock_model.side_effect = fake_inference + mock_load.return_value = mock_model + + frames = _make_test_frames(4) + result = interpolate_frames(frames, multiplier=2) + + # 4 frames with 2x: (4-1)*2 + 1 = 7 + assert len(result) == 7 + assert all(isinstance(f, Image.Image) for f in result) + + @patch("dw.tasks.interpolate_frames._load_rife_model") + def test_4x_quadruples_frame_count(self, mock_load): + """4x multiplier should run two passes of 2x.""" + from dw.tasks.interpolate_frames import interpolate_frames + + mock_model = MagicMock() + + def fake_inference(img1, img2): + arr1 = np.array(img1).astype(np.float32) + arr2 = np.array(img2).astype(np.float32) + mid = ((arr1 + arr2) / 2).astype(np.uint8) + return Image.fromarray(mid) + + mock_model.side_effect = fake_inference + mock_load.return_value = mock_model + + frames = _make_test_frames(4) + result = interpolate_frames(frames, multiplier=4) + + # Two passes of 2x: 4 -> 7 -> 13 + assert len(result) == 13 + assert all(isinstance(f, Image.Image) for f in result) + + def test_invalid_multiplier_raises(self): + """multiplier must be 2, 4, or 8.""" + from dw.tasks.interpolate_frames import interpolate_frames + + with pytest.raises(ValueError, match="multiplier"): + interpolate_frames(_make_test_frames(2), multiplier=3) + + def test_single_frame_raises(self): + """Need at least 2 frames to interpolate.""" + from dw.tasks.interpolate_frames import interpolate_frames + + with pytest.raises(ValueError, match="at least 2"): + interpolate_frames(_make_test_frames(1), multiplier=2) + + +class TestInterpolateFramesRegistration: + """Test that interpolate_frames is properly registered as a task command.""" + + def test_interpolate_frames_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + assert "interpolate_frames" in _COMMAND_REGISTRY +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/test_interpolate_frames.py -v` +Expected: FAIL — `dw.tasks.interpolate_frames` module does not exist yet. + +- [ ] **Step 3: Commit** + +```bash +git add tests/test_interpolate_frames.py +git commit -m "test: add interpolate_frames task tests" +``` + +--- + +### Task 7: Frame Interpolation Task — Implementation + +**Files:** +- Create: `dw/tasks/interpolate_frames.py` +- Modify: `dw/tasks/task.py:14` (add import), `dw/tasks/task.py` (add handler after segment handler) + +- [ ] **Step 1: Create dw/tasks/interpolate_frames.py** + +```python +""" +Video frame interpolation via RIFE (Real-Time Intermediate Flow Estimation). + +Takes a list of video frames and generates intermediate frames to increase +frame rate. Supports 2x, 4x, and 8x multipliers. + +Model weights are downloaded from HuggingFace Hub on first use. +""" + +import logging +import torch +import numpy as np +from PIL import Image + +logger = logging.getLogger("dw") + +_VALID_MULTIPLIERS = {2, 4, 8} + + +def interpolate_frames(video, device="cpu", **kwargs): + """Interpolate between video frames using RIFE to increase frame rate. + + Args: + video: List of PIL Images (video frames) + device: Target device ("cuda", "mps", "cpu") + **kwargs: + multiplier: Frame count multiplier — 2, 4, or 8 (default: 2) + model_name: HuggingFace repo with RIFE weights (default: auto) + + Returns: + List of PIL Images with interpolated frames inserted. + For N input frames with multiplier M, output is (N-1)*M + 1 frames + after log2(M) passes of 2x interpolation. + """ + multiplier = int(kwargs.get("multiplier", 2)) + model_name = kwargs.get("model_name", None) + + if multiplier not in _VALID_MULTIPLIERS: + raise ValueError( + f"multiplier must be one of {sorted(_VALID_MULTIPLIERS)}, got {multiplier}" + ) + + if len(video) < 2: + raise ValueError( + f"Need at least 2 frames to interpolate, got {len(video)}" + ) + + logger.info( + f"Interpolating {len(video)} frames with {multiplier}x multiplier on {device}" + ) + + model = _load_rife_model(device, model_name) + + # For 4x: two passes of 2x. For 8x: three passes of 2x. + passes = {2: 1, 4: 2, 8: 3}[multiplier] + frames = list(video) + + for pass_num in range(passes): + logger.debug( + f"Interpolation pass {pass_num + 1}/{passes}: {len(frames)} frames" + ) + frames = _interpolate_2x(frames, model) + + logger.info(f"Interpolation complete: {len(video)} -> {len(frames)} frames") + return frames + + +def _interpolate_2x(frames, model): + """Single pass of 2x interpolation — insert one frame between each pair.""" + result = [frames[0]] + for i in range(len(frames) - 1): + mid_frame = model(frames[i], frames[i + 1]) + result.append(mid_frame) + result.append(frames[i + 1]) + return result + + +def _load_rife_model(device, model_name=None): + """Load RIFE model and return a callable that interpolates two frames. + + Returns: + Callable that takes (frame1: PIL.Image, frame2: PIL.Image) -> PIL.Image + """ + try: + from huggingface_hub import hf_hub_download + except ImportError: + raise ImportError( + "huggingface_hub is required for RIFE model download. " + "Install with: pip install huggingface_hub" + ) + + if model_name is None: + model_name = "skytnt/anime-seg" # Placeholder — replace with actual RIFE HF repo + + logger.info(f"Loading RIFE model from {model_name} to {device}") + + # Download and load the RIFE IFNet model + # The exact loading code depends on the specific RIFE weights format on HF + # This is a functional wrapper pattern that abstracts the model internals + model_path = hf_hub_download(repo_id=model_name, filename="rife.pth") + + net = _build_ifnet() + state_dict = torch.load(model_path, map_location=device, weights_only=True) + net.load_state_dict(state_dict) + net.to(device) + + def inference(img1, img2): + """Interpolate a single frame between two input frames.""" + # Convert PIL to tensor [1, 3, H, W] in [0, 1] + arr1 = np.array(img1.convert("RGB")).astype(np.float32) / 255.0 + arr2 = np.array(img2.convert("RGB")).astype(np.float32) / 255.0 + t1 = torch.from_numpy(arr1).permute(2, 0, 1).unsqueeze(0).to(device) + t2 = torch.from_numpy(arr2).permute(2, 0, 1).unsqueeze(0).to(device) + + # Pad to multiple of 32 for RIFE + h, w = t1.shape[2], t1.shape[3] + ph = ((h - 1) // 32 + 1) * 32 + pw = ((w - 1) // 32 + 1) * 32 + padding = (0, pw - w, 0, ph - h) + t1_padded = torch.nn.functional.pad(t1, padding) + t2_padded = torch.nn.functional.pad(t2, padding) + + with torch.inference_mode(): + mid = net(t1_padded, t2_padded) + + # Remove padding and convert back to PIL + mid = mid[:, :, :h, :w] + mid = mid.squeeze(0).permute(1, 2, 0) + mid = (mid.clamp(0, 1) * 255).byte().cpu().numpy() + return Image.fromarray(mid) + + return inference + + +def _build_ifnet(): + """Build the RIFE IFNet architecture. + + NOTE: This is a placeholder. The actual IFNet architecture will need to be + vendored or adapted from the RIFE repository. The architecture is a + lightweight optical flow estimation network (~30MB weights). + + The real implementation should be adapted from: + https://github.com/hzwer/ECCV2022-RIFE + + For now this raises NotImplementedError to be replaced with the actual network. + """ + raise NotImplementedError( + "RIFE IFNet architecture needs to be vendored. " + "See https://github.com/hzwer/ECCV2022-RIFE for the source architecture." + ) +``` + +**Note:** The `_build_ifnet()` and `_load_rife_model()` functions contain the model loading scaffold. The actual RIFE IFNet architecture (~200 lines of PyTorch) will need to be vendored from the RIFE repository or adapted from a HuggingFace-hosted version. The test mocks bypass this by patching `_load_rife_model`. The frame list logic (`interpolate_frames`, `_interpolate_2x`) is fully functional. + +- [ ] **Step 2: Register the interpolate_frames command in task.py** + +Add import at the top of `dw/tasks/task.py`, after the `segment` import: + +```python +from .interpolate_frames import interpolate_frames +``` + +Add handler after the `_handle_segment` function: + +```python +@register_command("interpolate_frames") +def _handle_interpolate_frames(task, arguments, previous_pipelines): + """Interpolate video frames to increase frame rate""" + logger.debug("Interpolating frames") + video = arguments.pop("video") + return interpolate_frames(video, device=task.device, **arguments) +``` + +- [ ] **Step 3: Run tests to verify they pass** + +Run: `pytest tests/test_interpolate_frames.py -v` +Expected: All tests PASS (mocked model, so no downloads needed). + +- [ ] **Step 4: Run full test suite** + +Run: `pytest -v` +Expected: All tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add dw/tasks/interpolate_frames.py dw/tasks/task.py +git commit -m "feat: add interpolate_frames task (RIFE frame interpolation)" +``` + +--- + +### Task 8: Frame Interpolation Example Workflow + +**Files:** +- Create: `examples/InterpolateFrames.json` + +- [ ] **Step 1: Create examples/InterpolateFrames.json** + +```json +{ + "id": "InterpolateFrames", + "variables": { + "prompt": "A cat walking across a sunlit room", + "multiplier": 2 + }, + "steps": [ + { + "name": "generate_video", + "pipeline": { + "configuration": { + "offload": "sequential", + "component_type": "MochiPipeline", + "vae": { + "enable_tiling": true, + "enable_slicing": true + } + }, + "from_pretrained_arguments": { + "model_name": "genmo/mochi-1-preview", + "variant": "bf16", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "variable:prompt", + "num_frames": 85 + } + }, + "result": { + "content_type": "video/mp4", + "save": false, + "fps": 30 + } + }, + { + "name": "interpolate", + "task": { + "command": "interpolate_frames", + "arguments": { + "video": "previous_result:generate_video", + "multiplier": "variable:multiplier" + } + }, + "result": { + "content_type": "video/mp4", + "fps": 60 + } + } + ] +} +``` + +- [ ] **Step 2: Validate the workflow** + +Run: `python -m dw.validate examples/InterpolateFrames.json` +Expected: Validation passes. + +- [ ] **Step 3: Commit** + +```bash +git add examples/InterpolateFrames.json +git commit -m "feat: add InterpolateFrames example workflow" +``` + +--- + +### Task 9: Image Metadata Embedding — Tests + +**Files:** +- Modify: `tests/test_result.py` + +- [ ] **Step 1: Add metadata embedding tests to test_result.py** + +Add these test classes at the end of the file, before the `if __name__` block: + +```python +class TestMetadataEmbedding: + """Test opt-in metadata embedding in saved images.""" + + def test_png_metadata_embedded(self): + """When embed_metadata is true, PNG should contain parameters text chunk.""" + with tempfile.TemporaryDirectory() as temp_dir: + result_def = {"content_type": "image/png", "save": True, "embed_metadata": True} + result = Result(result_def) + result.set_metadata({ + "workflow_id": "test_workflow", + "step_name": "generate", + "model_name": "test/model", + "arguments": {"prompt": "a cat", "num_inference_steps": 25}, + }) + + # Add a real PIL image + img = Image.new("RGB", (64, 64), color=(128, 64, 32)) + result.add_result(img) + result.save(temp_dir, "test_output") + + # Read back the PNG and check for metadata + output_file = os.path.join(temp_dir, "test_output-0.0.png") + assert os.path.exists(output_file) + + saved_img = Image.open(output_file) + assert "parameters" in saved_img.info + metadata = json.loads(saved_img.info["parameters"]) + assert metadata["workflow_id"] == "test_workflow" + assert metadata["arguments"]["prompt"] == "a cat" + + def test_no_metadata_when_not_enabled(self): + """When embed_metadata is absent/false, no metadata should be embedded.""" + with tempfile.TemporaryDirectory() as temp_dir: + result_def = {"content_type": "image/png", "save": True} + result = Result(result_def) + + img = Image.new("RGB", (64, 64), color=(128, 64, 32)) + result.add_result(img) + result.save(temp_dir, "test_output") + + output_file = os.path.join(temp_dir, "test_output-0.0.png") + saved_img = Image.open(output_file) + assert "parameters" not in saved_img.info + + def test_set_metadata_method(self): + """set_metadata should store metadata on the Result instance.""" + result = Result({}) + assert result.metadata is None + + metadata = {"workflow_id": "test", "step_name": "step1"} + result.set_metadata(metadata) + assert result.metadata == metadata + + def test_metadata_with_embed_false(self): + """When embed_metadata is explicitly false, no metadata embedded even if set.""" + with tempfile.TemporaryDirectory() as temp_dir: + result_def = { + "content_type": "image/png", + "save": True, + "embed_metadata": False, + } + result = Result(result_def) + result.set_metadata({"workflow_id": "test"}) + + img = Image.new("RGB", (64, 64), color=(128, 64, 32)) + result.add_result(img) + result.save(temp_dir, "test_output") + + output_file = os.path.join(temp_dir, "test_output-0.0.png") + saved_img = Image.open(output_file) + assert "parameters" not in saved_img.info +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/test_result.py::TestMetadataEmbedding -v` +Expected: FAIL — `Result` has no `set_metadata` method or `metadata` attribute yet. + +- [ ] **Step 3: Commit** + +```bash +git add tests/test_result.py +git commit -m "test: add metadata embedding tests" +``` + +--- + +### Task 10: Image Metadata Embedding — Result Implementation + +**Files:** +- Modify: `dw/result.py:24-35` (Result.__init__), `dw/result.py:159-238` (save_artifact) + +- [ ] **Step 1: Add metadata support to Result class** + +In `dw/result.py`, add `self.metadata = None` to `__init__` (after line 34): + +```python + def __init__(self, result_definition): + self.result_definition = result_definition + self.result_list = [] + self.metadata = None + logger.debug(f"Initialized Result with definition: {result_definition}") +``` + +Add `set_metadata` method after `add_result` (after line 51): + +```python + def set_metadata(self, metadata): + """Set metadata to embed in saved image files. + + Args: + metadata: Dict of generation parameters to embed + """ + self.metadata = metadata + logger.debug(f"Set metadata for result: {list(metadata.keys())}") +``` + +- [ ] **Step 2: Modify save_artifact to embed PNG metadata** + +In `dw/result.py`, modify the image saving branch in `save_artifact`. The current code at line 228 is: + +```python + elif hasattr(artifact, "save"): + artifact.save(output_path) +``` + +Replace it with metadata-aware saving: + +```python + elif hasattr(artifact, "save"): + if ( + self.metadata is not None + and self.result_definition.get("embed_metadata", False) + and content_type.startswith("image/") + ): + self._save_image_with_metadata(artifact, output_path, content_type) + else: + artifact.save(output_path) +``` + +Add the `_save_image_with_metadata` method to the Result class (after `save_artifact`): + +```python + def _save_image_with_metadata(self, image, output_path, content_type): + """Save an image with embedded generation metadata. + + Args: + image: PIL Image to save + output_path: File path to save to + content_type: MIME type (determines embedding method) + """ + metadata_json = json.dumps(self.metadata, default=str) + + if content_type == "image/png": + from PIL.PngImagePlugin import PngInfo + + png_info = PngInfo() + png_info.add_text("parameters", metadata_json) + image.save(output_path, pnginfo=png_info) + logger.debug(f"Embedded PNG metadata in {output_path}") + elif content_type in ("image/jpeg", "image/webp"): + try: + import piexif + + exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}} + if hasattr(image, "info") and "exif" in image.info: + exif_dict = piexif.load(image.info["exif"]) + exif_dict["Exif"][ + piexif.ExifIFD.UserComment + ] = piexif.helper.UserComment.dump(metadata_json) + exif_bytes = piexif.dump(exif_dict) + image.save(output_path, exif=exif_bytes) + logger.debug(f"Embedded EXIF metadata in {output_path}") + except ImportError: + logger.warning( + "piexif not installed - saving without metadata. " + "Install with: pip install piexif" + ) + image.save(output_path) + else: + # Unsupported image format for metadata - save normally + image.save(output_path) +``` + +- [ ] **Step 3: Run metadata tests** + +Run: `pytest tests/test_result.py::TestMetadataEmbedding -v` +Expected: All 4 tests PASS. + +- [ ] **Step 4: Run full test suite** + +Run: `pytest -v` +Expected: All tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add dw/result.py +git commit -m "feat: add opt-in image metadata embedding to Result" +``` + +--- + +### Task 11: Image Metadata Embedding — Step Integration + +**Files:** +- Modify: `dw/step.py:26-96` +- Modify: `dw/workflow_schema.json:650-689` + +- [ ] **Step 1: Update workflow schema to document embed_metadata** + +In `dw/workflow_schema.json`, add `embed_metadata` to the result definition properties (after the `samplerate` property, before the closing `}` of properties around line 675): + +```json + "embed_metadata": { + "description": "Whether to embed generation parameters as metadata in saved images (PNG info chunks or EXIF). Only applies to image content types.", + "type": "boolean", + "default": false + } +``` + +- [ ] **Step 2: Modify step.py to pass metadata to Result** + +In `dw/step.py`, modify the `run` method. After the line that creates the Result (line 41) add metadata collection: + +```python + result = Result(self.step_definition.get("result", {})) + + # Collect metadata for embedding if enabled + result_def = self.step_definition.get("result", {}) + if result_def.get("embed_metadata", False): + metadata = self._collect_metadata() + result.set_metadata(metadata) +``` + +Add the `_collect_metadata` method to the Step class (after the `run` method): + +```python + def _collect_metadata(self): + """Collect step metadata for embedding in saved images. + + Returns: + Dict of generation parameters from this step. + """ + metadata = {"step_name": self.name} + + if "pipeline" in self.step_definition: + pipeline_def = self.step_definition["pipeline"] + pretrained_args = pipeline_def.get("from_pretrained_arguments", {}) + if "model_name" in pretrained_args: + metadata["model_name"] = pretrained_args["model_name"] + + # Copy inference arguments (prompt, steps, guidance, etc.) + metadata["arguments"] = dict(pipeline_def.get("arguments", {})) + + elif "task" in self.step_definition: + task_def = self.step_definition["task"] + metadata["task_command"] = task_def.get("command", "unknown") + metadata["arguments"] = dict(task_def.get("arguments", {})) + + return metadata +``` + +- [ ] **Step 3: Run all tests** + +Run: `pytest -v` +Expected: All tests pass. + +- [ ] **Step 4: Commit** + +```bash +git add dw/step.py dw/workflow_schema.json +git commit -m "feat: wire metadata collection from step to result" +``` + +--- + +### Task 12: Metadata Embedding Example Workflow + +**Files:** +- Create: `examples/MetadataEmbed.json` + +- [ ] **Step 1: Create examples/MetadataEmbed.json** + +```json +{ + "id": "MetadataEmbed", + "variables": { + "prompt": "A serene mountain landscape at golden hour, photorealistic", + "steps": 25 + }, + "steps": [ + { + "name": "generate", + "pipeline": { + "configuration": { + "offload": "model", + "component_type": "FluxPipeline" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "variable:prompt", + "num_inference_steps": "variable:steps", + "guidance_scale": 3.5 + } + }, + "result": { + "content_type": "image/png", + "embed_metadata": true + } + } + ] +} +``` + +- [ ] **Step 2: Validate the workflow** + +Run: `python -m dw.validate examples/MetadataEmbed.json` +Expected: Validation passes. + +- [ ] **Step 3: Commit** + +```bash +git add examples/MetadataEmbed.json +git commit -m "feat: add MetadataEmbed example workflow" +``` + +--- + +### Task 13: Final Verification + +- [ ] **Step 1: Run the full test suite** + +Run: `pytest -v` +Expected: All tests pass, including new segment, interpolation, and metadata tests. + +- [ ] **Step 2: Validate all new example workflows** + +Run: `python -m dw.validate examples/MarigoldDepth.json && python -m dw.validate examples/MarigoldNormals.json && python -m dw.validate examples/Segment.json && python -m dw.validate examples/SegmentAndInpaint.json && python -m dw.validate examples/InterpolateFrames.json && python -m dw.validate examples/MetadataEmbed.json` +Expected: All 6 pass validation. + +- [ ] **Step 3: Run black formatting** + +Run: `black dw/ tests/` +Expected: Files formatted (or already formatted). + +- [ ] **Step 4: Final commit if formatting changed anything** + +```bash +git add -A +git commit -m "style: format new files with black" +``` diff --git a/docs/superpowers/specs/2026-03-26-ecosystem-utilities-design.md b/docs/superpowers/specs/2026-03-26-ecosystem-utilities-design.md new file mode 100644 index 00000000..1c28fb0d --- /dev/null +++ b/docs/superpowers/specs/2026-03-26-ecosystem-utilities-design.md @@ -0,0 +1,363 @@ +# Ecosystem Utilities Design Spec + +Four new capabilities from the stable diffusion ecosystem: Marigold depth/normals examples, GroundingDINO+SAM2 segmentation task, RIFE frame interpolation task, and image metadata embedding. + +## 1. Marigold Depth & Normals (Example Workflows Only) + +Marigold ships as native diffusers pipelines (`MarigoldDepthPipeline`, `MarigoldNormalsPipeline`). No code changes needed — just example workflows. + +### Files to Create + +- `examples/MarigoldDepth.json` +- `examples/MarigoldNormals.json` + +### MarigoldDepth.json + +Two-step workflow: +1. `gather_images` — load an input image (URL or local path via variable) +2. Pipeline step using `MarigoldDepthPipeline` with `prs-eth/marigold-depth-lcm-v1-0` — outputs a depth map as a PIL Image + +The depth output chains naturally into ControlNet steps via `previous_result:`. + +### MarigoldNormals.json + +Same structure using `MarigoldNormalsPipeline` with `prs-eth/marigold-normals-lcm-v1-0`. + +### Dependencies + +None beyond existing diffusers install. + +--- + +## 2. GroundingDINO + SAM2 Segmentation Task + +A new `segment` task command that takes an image + text prompt and returns a binary mask image. + +### Files to Create/Modify + +- `dw/tasks/segment.py` — core implementation +- `dw/tasks/task.py` — import and register command +- `examples/Segment.json` — basic segmentation example +- `examples/SegmentAndInpaint.json` — segment then inpaint with FluxFill + +### Task Registration + +```python +# In task.py +from .segment import segment_image + +@register_command("segment") +def _handle_segment(task, arguments, previous_pipelines): + """Segment objects in an image using text prompt""" + logger.debug("Segmenting image") + image = arguments.pop("image") + prompt = arguments.pop("prompt") + return segment_image(image, prompt, device=task.device, **arguments) +``` + +### Task Arguments + +| Argument | Type | Required | Default | Description | +|----------|------|----------|---------|-------------| +| `image` | PIL Image | yes | — | Input image | +| `prompt` | string | yes | — | Text description of object(s) to segment | +| `model_name` | string | no | `"IDEA-Research/grounding-dino-base"` | GroundingDINO model | +| `sam_model_name` | string | no | `"facebook/sam2-hiera-large"` | SAM2 model | +| `threshold` | float | no | `0.3` | Detection confidence threshold | +| `invert` | bool | no | `false` | Invert the output mask | + +### Returns + +PIL Image — binary mask (white = detected object, black = background). When multiple objects match the prompt, masks are combined via union. + +### Implementation: `dw/tasks/segment.py` + +```python +def segment_image(image, prompt, device="cpu", **kwargs): + """Segment objects matching text prompt, returning a binary mask.""" +``` + +Flow: +1. Load GroundingDINO via `transformers.AutoProcessor` + `AutoModelForZeroShotObjectDetection` +2. Run detection with text prompt, filter by `threshold` +3. Extract bounding boxes +4. Load SAM2 via `transformers.AutoProcessor` + `AutoModelForMaskGeneration` (or `SamModel` + `SamProcessor`) +5. Load SAM2 via `Sam2Model.from_pretrained()` + `Sam2Processor.from_pretrained()` from `transformers` +6. Feed image + bounding boxes to SAM2 +7. Combine all output masks (union) into single binary mask +8. Optionally invert +9. Return as PIL Image (mode "L", 0/255 values) + +Direct top-level imports for `transformers` model classes (`AutoModelForZeroShotObjectDetection`, `AutoProcessor`, `Sam2Model`, `Sam2Processor`) since `transformers` is a hard dependency of this project. + +### Dependencies + +- `transformers` (already likely installed) +- `torch` (already installed) + +### Example: Segment.json + +```json +{ + "id": "Segment", + "variables": { + "image_url": "https://example.com/photo.jpg", + "prompt": "dog" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "segment", + "task": { + "command": "segment", + "arguments": { + "image": "previous_result:load_image", + "prompt": "variable:prompt" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} +``` + +### Example: SegmentAndInpaint.json + +Three-step workflow: +1. `gather_images` — load input image +2. `segment` — generate mask from text prompt +3. `FluxFillPipeline` step — inpaint using mask from step 2 and a new prompt + +--- + +## 3. Frame Interpolation Task (RIFE) + +A new `interpolate_frames` task that takes video frames and returns interpolated frames at higher frame count. + +### Files to Create/Modify + +- `dw/tasks/interpolate_frames.py` — core implementation +- `dw/tasks/task.py` — import and register command +- `examples/InterpolateFrames.json` — video gen + interpolation example + +### Task Registration + +```python +# In task.py +from .interpolate_frames import interpolate_frames + +@register_command("interpolate_frames") +def _handle_interpolate_frames(task, arguments, previous_pipelines): + """Interpolate video frames to increase frame rate""" + logger.debug("Interpolating frames") + video = arguments.pop("video") + return interpolate_frames(video, device=task.device, **arguments) +``` + +### Task Arguments + +| Argument | Type | Required | Default | Description | +|----------|------|----------|---------|-------------| +| `video` | list[PIL Image] | yes | — | Input video frames | +| `multiplier` | int | no | `2` | Frame count multiplier (2, 4, or 8) | +| `model_name` | string | no | see below | RIFE model weights (HF repo or local path) | + +### Returns + +List of PIL Images — the interpolated frame sequence. The caller should set `fps` in the result config to `original_fps * multiplier` to maintain correct playback speed. + +### Implementation: `dw/tasks/interpolate_frames.py` + +```python +def interpolate_frames(video, device="cpu", **kwargs): + """Interpolate between video frames using RIFE.""" +``` + +Flow: +1. Validate `multiplier` is 2, 4, or 8 +2. Load RIFE model (PyTorch implementation for device compatibility) +3. For each consecutive frame pair, generate intermediate frame(s) +4. For 4x: run two passes of 2x interpolation +5. For 8x: run three passes of 2x interpolation +6. Return combined frame list + +Uses a vendored PyTorch RIFE IFNet architecture with weights downloaded from HuggingFace Hub at first use. This ensures broad device compatibility across CUDA, MPS, and CPU without requiring ncnn or Vulkan. The IFNet model files are small (~30MB) and cached locally after first download. Lazy imports. + +### Dependencies + +- `torch` (already installed) +- RIFE model weights downloaded from HuggingFace Hub at first use + +### Example: InterpolateFrames.json + +```json +{ + "id": "InterpolateFrames", + "variables": { + "prompt": "A cat walking across a sunlit room", + "multiplier": 2 + }, + "steps": [ + { + "name": "generate_video", + "pipeline": { + "configuration": { + "offload": "sequential", + "component_type": "MochiPipeline", + "vae": { "enable_tiling": true, "enable_slicing": true } + }, + "from_pretrained_arguments": { + "model_name": "genmo/mochi-1-preview", + "variant": "bf16", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "variable:prompt", + "num_frames": 85 + } + }, + "result": { + "content_type": "video/mp4", + "save": false, + "fps": 30 + } + }, + { + "name": "interpolate", + "task": { + "command": "interpolate_frames", + "arguments": { + "video": "previous_result:generate_video", + "multiplier": "variable:multiplier" + } + }, + "result": { + "content_type": "video/mp4", + "fps": 60 + } + } + ] +} +``` + +--- + +## 4. Image Metadata Embedding + +Opt-in embedding of generation parameters into saved image files via the `result` config. + +### Files to Modify + +- `dw/workflow_schema.json` — add `embed_metadata` property to result definition +- `dw/result.py` — embed metadata when saving images +- `dw/step.py` — pass step metadata to Result when `embed_metadata` is true +- `examples/MetadataEmbed.json` — example workflow + +### Schema Change + +Add to the `result` definition properties in `workflow_schema.json`: + +```json +"embed_metadata": { + "description": "Whether to embed generation parameters as metadata in saved images (PNG info chunks or EXIF)", + "type": "boolean", + "default": false +} +``` + +### What Gets Embedded + +A JSON object stored as a text chunk containing: +- `workflow_id` — the workflow's ID +- `step_name` — the step that produced the image +- `model_name` — from `from_pretrained_arguments` (if pipeline step) +- `arguments` — the step's inference arguments (prompt, negative_prompt, guidance_scale, num_inference_steps, seed, etc.) +- `task_command` — if a task step, the command name + +### Implementation Changes + +**`step.py`:** When `embed_metadata` is true in the step's result definition, collect metadata and pass it to `Result` via a new `set_metadata(metadata_dict)` method. Metadata is collected from: + +- `step_definition["name"]` — the step name +- The workflow ID (passed through from `workflow.py`) +- For pipeline steps: `from_pretrained_arguments.model_name`, and the inference `arguments` dict (prompt, negative_prompt, guidance_scale, num_inference_steps, seed, width, height, etc.) +- For task steps: `task.command` and the task `arguments` dict + +**`result.py`:** +- Add `self.metadata = None` to `Result.__init__` +- Add `set_metadata(self, metadata)` method +- In `save_artifact`, when `content_type` starts with `image/` and `self.metadata` is not None: + - **PNG:** Use `PIL.PngImagePlugin.PngInfo` to add a `parameters` text chunk containing the JSON metadata. Pass `pnginfo` to `image.save()`. + - **JPEG/WebP:** Use `piexif` (lazy import, optional) to write metadata as EXIF UserComment. If `piexif` is not installed, log a warning and save without metadata. + +### Format + +PNG text chunk key: `parameters` +Value: JSON string of the metadata dict. + +This is compatible with tools that read A1111-style PNG info, though the structure is JSON rather than A1111's custom text format. + +### Example: MetadataEmbed.json + +```json +{ + "id": "MetadataEmbed", + "variables": { + "prompt": "A serene mountain landscape at golden hour", + "steps": 25 + }, + "steps": [ + { + "name": "generate", + "pipeline": { + "configuration": { + "offload": "model", + "component_type": "FluxPipeline" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "variable:prompt", + "num_inference_steps": "variable:steps", + "guidance_scale": 3.5 + } + }, + "result": { + "content_type": "image/png", + "embed_metadata": true + } + } + ] +} +``` + +### Dependencies + +- `piexif` — optional, only needed for JPEG/WebP metadata. PNG works with PIL alone. + +--- + +## Summary of Changes + +| Feature | New Files | Modified Files | New Dependencies | +|---------|-----------|----------------|------------------| +| Marigold examples | 2 example JSONs | none | none | +| Segment task | `dw/tasks/segment.py`, 2 example JSONs | `dw/tasks/task.py` | `transformers` (likely already present) | +| Frame interpolation | `dw/tasks/interpolate_frames.py`, 1 example JSON | `dw/tasks/task.py` | RIFE model weights (auto-download) | +| Metadata embedding | 1 example JSON | `dw/result.py`, `dw/step.py`, `dw/workflow_schema.json` | `piexif` (optional, JPEG only) | diff --git a/dw/result.py b/dw/result.py index 8a9b439e..e55dc6a8 100644 --- a/dw/result.py +++ b/dw/result.py @@ -32,8 +32,17 @@ def __init__(self, result_definition): """ self.result_definition = result_definition self.result_list = [] + self.metadata = None logger.debug(f"Initialized Result with definition: {result_definition}") + def set_metadata(self, metadata): + """Set metadata to embed in saved image artifacts. + + Args: + metadata: Dict of generation parameters to embed + """ + self.metadata = metadata + def add_result(self, result): """Add one or more results to the result list. @@ -226,7 +235,14 @@ def save_artifact( with open(output_path, "w") as file: file.write(artifact) elif hasattr(artifact, "save"): - artifact.save(output_path) + if ( + self.metadata is not None + and self.result_definition.get("embed_metadata", False) + and content_type.startswith("image/") + ): + self._save_image_with_metadata(artifact, output_path, content_type) + else: + artifact.save(output_path) else: raise ValueError( f"Content type {content_type} does not match result type {type(artifact)}" @@ -237,6 +253,45 @@ def save_artifact( ) raise + def _save_image_with_metadata(self, image, output_path, content_type): + """Save an image with embedded generation metadata. + + Args: + image: PIL Image to save + output_path: Path to save the image to + content_type: MIME type of the image + """ + metadata_json = json.dumps(self.metadata, default=str) + + if content_type == "image/png": + from PIL.PngImagePlugin import PngInfo + + png_info = PngInfo() + png_info.add_text("parameters", metadata_json) + image.save(output_path, pnginfo=png_info) + logger.debug(f"Embedded PNG metadata in {output_path}") + elif content_type in ("image/jpeg", "image/webp"): + try: + import piexif + + exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}} + if hasattr(image, "info") and "exif" in image.info: + exif_dict = piexif.load(image.info["exif"]) + exif_dict["Exif"][piexif.ExifIFD.UserComment] = ( + piexif.helper.UserComment.dump(metadata_json) + ) + exif_bytes = piexif.dump(exif_dict) + image.save(output_path, exif=exif_bytes) + logger.debug(f"Embedded EXIF metadata in {output_path}") + except ImportError: + logger.warning( + "piexif not installed - saving without metadata. " + "Install with: pip install piexif" + ) + image.save(output_path) + else: + image.save(output_path) + def get_artifact_list(result): """Extract list of artifacts from a result object. diff --git a/dw/step.py b/dw/step.py index 496a2acf..a7d90c92 100644 --- a/dw/step.py +++ b/dw/step.py @@ -40,6 +40,12 @@ def run(self, previous_results, previous_pipelines, step_action): # This handles how results should be saved/processed result = Result(self.step_definition.get("result", {})) + # Collect metadata for embedding if enabled + result_def = self.step_definition.get("result", {}) + if result_def.get("embed_metadata", False): + metadata = self._collect_metadata() + result.set_metadata(metadata) + # Log what type of action we're executing (Pipeline/Task/Workflow) action_type = type(step_action).__name__ logger.info(f"Running {action_type} {step_name}:{step_action.name}...") @@ -114,3 +120,21 @@ def run(self, previous_results, previous_pipelines, step_action): exc_info=True, ) raise + + def _collect_metadata(self): + """Collect step metadata for embedding in saved images.""" + metadata = {"step_name": self.name} + + if "pipeline" in self.step_definition: + pipeline_def = self.step_definition["pipeline"] + pretrained_args = pipeline_def.get("from_pretrained_arguments", {}) + if "model_name" in pretrained_args: + metadata["model_name"] = pretrained_args["model_name"] + metadata["arguments"] = dict(pipeline_def.get("arguments", {})) + + elif "task" in self.step_definition: + task_def = self.step_definition["task"] + metadata["task_command"] = task_def.get("command", "unknown") + metadata["arguments"] = dict(task_def.get("arguments", {})) + + return metadata diff --git a/dw/tasks/diffusion_upscale.py b/dw/tasks/diffusion_upscale.py new file mode 100644 index 00000000..de297dc4 --- /dev/null +++ b/dw/tasks/diffusion_upscale.py @@ -0,0 +1,96 @@ +""" +Diffusion-based image upscaling via Stable Diffusion upscale pipelines. + +Provides text-guided upscaling with better detail recovery than +traditional super-resolution models, especially for faces and textures. + +Supports two modes: + - "x4" (default): StableDiffusionUpscalePipeline (4x, stabilityai/stable-diffusion-x4-upscaler) + - "x2": StableDiffusionLatentUpscalePipeline (2x, stabilityai/sd-x2-latent-upscaler) +""" + +import logging +import torch +import diffusers + +logger = logging.getLogger("dw") + +_MODELS = { + "x4": { + "pipeline_class": "StableDiffusionUpscalePipeline", + "model_name": "stabilityai/stable-diffusion-x4-upscaler", + }, + "x2": { + "pipeline_class": "StableDiffusionLatentUpscalePipeline", + "model_name": "stabilityai/sd-x2-latent-upscaler", + }, +} + + +def diffusion_upscale(image, device="cpu", **kwargs): + """Upscale an image using a Stable Diffusion upscale pipeline. + + Args: + image: PIL Image to upscale. + device: Target device ("cuda", "mps", "cpu"). + **kwargs: + prompt: Text guidance for upscaling (default: ""). + negative_prompt: Negative text guidance (default: None). + mode: "x4" or "x2" (default: "x4"). + model_name: Override the default model for the selected mode. + num_inference_steps: Denoising steps (default: 25). + guidance_scale: Classifier-free guidance scale (default: 9.0). + noise_level: Noise level for x4 mode (default: 20, ignored for x2). + + Returns: + PIL Image (upscaled). + """ + mode = kwargs.get("mode", "x4") + if mode not in _MODELS: + raise ValueError(f"mode must be one of {sorted(_MODELS.keys())}, got '{mode}'") + + config = _MODELS[mode] + model_name = kwargs.get("model_name", config["model_name"]) + prompt = kwargs.get("prompt", "") + negative_prompt = kwargs.get("negative_prompt", None) + num_inference_steps = int(kwargs.get("num_inference_steps", 25)) + guidance_scale = float(kwargs.get("guidance_scale", 9.0)) + noise_level = int(kwargs.get("noise_level", 20)) + + pipeline_class = getattr(diffusers, config["pipeline_class"]) + + logger.info( + f"Loading {config['pipeline_class']} from {model_name} to {device}" + ) + + dtype = torch.float16 if device == "cuda" else torch.float32 + pipe = pipeline_class.from_pretrained( + model_name, + torch_dtype=dtype, + ) + pipe.to(device) + + call_kwargs = { + "prompt": prompt, + "image": image, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + } + + if negative_prompt is not None: + call_kwargs["negative_prompt"] = negative_prompt + + if mode == "x4": + call_kwargs["noise_level"] = noise_level + + logger.info( + f"Upscaling {image.width}x{image.height} with {mode} mode, " + f"{num_inference_steps} steps" + ) + + with torch.inference_mode(): + result = pipe(**call_kwargs) + + output = result.images[0] + logger.info(f"Upscaled to {output.width}x{output.height}") + return output diff --git a/dw/tasks/image_to_text.py b/dw/tasks/image_to_text.py new file mode 100644 index 00000000..e1d3543c --- /dev/null +++ b/dw/tasks/image_to_text.py @@ -0,0 +1,55 @@ +""" +Image-to-text captioning via HuggingFace transformers pipeline. + +Supports BLIP, BLIP-2, ViT-GPT2, GIT, and other models compatible +with the transformers image-to-text pipeline. +""" + +import logging +import torch +from transformers import pipeline as hf_pipeline + +logger = logging.getLogger("dw") + +_DEFAULT_MODEL = "Salesforce/blip-image-captioning-base" + + +def image_to_text(image, device="cpu", **kwargs): + """Generate a text caption for an image. + + Args: + image: PIL Image to caption. + device: Target device ("cuda", "mps", "cpu"). + **kwargs: + model_name: HuggingFace model ID (default: Salesforce/blip-image-captioning-base). + prompt: Optional text prompt for conditional captioning (BLIP-2, etc.). + max_new_tokens: Max tokens to generate (default: 50). + + Returns: + Caption string. + """ + model_name = kwargs.get("model_name", _DEFAULT_MODEL) + prompt = kwargs.get("prompt", None) + max_new_tokens = int(kwargs.get("max_new_tokens", 50)) + + logger.info(f"Captioning image with {model_name} on {device}") + + dtype = torch.float16 if device == "cuda" else torch.float32 + # Use device_map instead of device to avoid caching_allocator_warmup + # buffer pre-allocation failures on MPS and with larger models. + pipe = hf_pipeline( + "image-to-text", + model=model_name, + device_map=device, + torch_dtype=dtype, + ) + + generate_kwargs = {"max_new_tokens": max_new_tokens} + if prompt is not None: + generate_kwargs["prompt"] = prompt + + results = pipe(image, generate_kwargs=generate_kwargs) + + caption = results[0]["generated_text"].strip() + logger.info(f"Caption: {caption[:100]}{'...' if len(caption) > 100 else ''}") + return caption diff --git a/dw/tasks/image_utils.py b/dw/tasks/image_utils.py index eee0e837..2e7f835a 100644 --- a/dw/tasks/image_utils.py +++ b/dw/tasks/image_utils.py @@ -158,6 +158,15 @@ def process_image(image, processor, device, kwargs): if processor == "resize_rescale": return resize_rescale(image, **kwargs) + if processor == "resize_bucket": + return resize_bucket(image, **kwargs) + + if processor == "strip_exif": + return strip_exif(image) + + if processor == "add_watermark": + return add_watermark(image, **kwargs) + raise Exception(f"Unknown image processor type: {processor}") @@ -300,6 +309,137 @@ def resize_resample(image, resolution=1024): return input_image.resize((W, H), resample=Image.LANCZOS) +# Standard aspect ratios used by SDXL, Flux, and similar models. +# Each entry is (width_ratio, height_ratio). +_DEFAULT_RATIOS = [ + (1, 1), + (4, 3), + (3, 4), + (3, 2), + (2, 3), + (16, 9), + (9, 16), + (21, 9), + (9, 21), +] + + +def resize_bucket(image, resolution=1024, ratios=None, alignment=64): + """Resize image to the closest model-native aspect ratio bucket. + + Picks the standard ratio closest to the input image's natural aspect + ratio, then scales to fit within the target resolution (based on the + short side) with dimensions aligned to `alignment` pixels. + + Args: + image: PIL Image to resize. + resolution: Target size for the short side in pixels (default: 1024). + ratios: Optional list of [w, h] ratio pairs. Defaults to standard + ratios used by SDXL/Flux (1:1, 4:3, 3:2, 16:9, etc.). + alignment: Round dimensions to this multiple (default: 64). + + Returns: + PIL Image resized to the bucketed dimensions. + """ + input_image = image.convert("RGB") + W, H = input_image.size + input_ratio = W / H + + bucket_ratios = ratios if ratios is not None else _DEFAULT_RATIOS + + # Find the closest aspect ratio + best_ratio = min( + bucket_ratios, + key=lambda r: abs((r[0] / r[1]) - input_ratio), + ) + + wr, hr = best_ratio + bucket_ratio = wr / hr + + # Scale so the short side matches resolution, then align + if bucket_ratio >= 1.0: + # Landscape or square: height is the short side + out_h = int(round(resolution / alignment)) * alignment + out_w = int(round((out_h * bucket_ratio) / alignment)) * alignment + else: + # Portrait: width is the short side + out_w = int(round(resolution / alignment)) * alignment + out_h = int(round((out_w / bucket_ratio) / alignment)) * alignment + + return input_image.resize((out_w, out_h), resample=Image.LANCZOS) + + +def strip_exif(image): + """Remove all EXIF and metadata from an image. + + Creates a clean copy with pixel data only — no GPS coordinates, + camera info, timestamps, or other embedded metadata. + + Args: + image: PIL Image to strip. + + Returns: + PIL Image with all metadata removed. + """ + clean = Image.new(image.mode, image.size) + clean.paste(image) + return clean + + +def add_watermark(image, text="AI Generated", position="bottom-right", + opacity=128, font_size=0, margin=10, color=None): + """Add a visible text watermark to an image. + + Args: + image: PIL Image to watermark. + text: Watermark text (default: "AI Generated"). + position: Placement — "bottom-right", "bottom-left", "top-right", + "top-left", or "center" (default: "bottom-right"). + opacity: Text opacity 0-255 (default: 128). + font_size: Font size in pixels. 0 = auto-scale to ~3% of image height. + margin: Pixel margin from edges (default: 10). + color: RGB tuple for text color (default: white). + + Returns: + PIL Image with watermark applied. + """ + from PIL import ImageDraw, ImageFont + + base = image.convert("RGBA") + overlay = Image.new("RGBA", base.size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay) + + if color is None: + color = (255, 255, 255) + fill = (*color, int(opacity)) + + if font_size <= 0: + font_size = max(12, base.height // 30) + + try: + font = ImageFont.truetype("Arial", font_size) + except (IOError, OSError): + font = ImageFont.load_default(size=font_size) + + bbox = draw.textbbox((0, 0), text, font=font) + text_w = bbox[2] - bbox[0] + text_h = bbox[3] - bbox[1] + + positions = { + "bottom-right": (base.width - text_w - margin, base.height - text_h - margin), + "bottom-left": (margin, base.height - text_h - margin), + "top-right": (base.width - text_w - margin, margin), + "top-left": (margin, margin), + "center": ((base.width - text_w) // 2, (base.height - text_h) // 2), + } + xy = positions.get(position, positions["bottom-right"]) + + draw.text(xy, text, font=font, fill=fill) + + result = Image.alpha_composite(base, overlay) + return result.convert("RGB") + + ada_palette = np.asarray( [ [0, 0, 0], diff --git a/dw/tasks/interpolate_frames.py b/dw/tasks/interpolate_frames.py new file mode 100644 index 00000000..10ba089b --- /dev/null +++ b/dw/tasks/interpolate_frames.py @@ -0,0 +1,153 @@ +""" +Video frame interpolation via RIFE (Real-Time Intermediate Flow Estimation). + +Takes a list of video frames and generates intermediate frames to increase +frame rate. Supports 2x, 4x, and 8x multipliers. + +Model weights are downloaded from HuggingFace Hub on first use. +""" + +import logging +import torch +import numpy as np +from PIL import Image + +logger = logging.getLogger("dw") + +_VALID_MULTIPLIERS = {2, 4, 8} + + +def interpolate_frames(video, device="cpu", **kwargs): + """Interpolate between video frames using RIFE to increase frame rate. + + Args: + video: List of PIL Images (video frames) + device: Target device ("cuda", "mps", "cpu") + **kwargs: + multiplier: Frame count multiplier — 2, 4, or 8 (default: 2) + model_name: HuggingFace repo with RIFE weights (default: auto) + + Returns: + List of PIL Images with interpolated frames inserted. + """ + multiplier = int(kwargs.get("multiplier", 2)) + model_name = kwargs.get("model_name", None) + + if multiplier not in _VALID_MULTIPLIERS: + raise ValueError( + f"multiplier must be one of {sorted(_VALID_MULTIPLIERS)}, got {multiplier}" + ) + + if len(video) < 2: + raise ValueError(f"Need at least 2 frames to interpolate, got {len(video)}") + + logger.info( + f"Interpolating {len(video)} frames with {multiplier}x multiplier on {device}" + ) + + model = _load_rife_model(device, model_name) + + passes = {2: 1, 4: 2, 8: 3}[multiplier] + frames = list(video) + + for pass_num in range(passes): + logger.debug( + f"Interpolation pass {pass_num + 1}/{passes}: {len(frames)} frames" + ) + frames = _interpolate_2x(frames, model) + + logger.info(f"Interpolation complete: {len(video)} -> {len(frames)} frames") + return frames + + +def _interpolate_2x(frames, model): + """Single pass of 2x interpolation — insert one frame between each pair.""" + result = [frames[0]] + for i in range(len(frames) - 1): + mid_frame = model(frames[i], frames[i + 1]) + result.append(mid_frame) + result.append(frames[i + 1]) + return result + + +_DEFAULT_RIFE_REPO = "styler00dollar/RIFE-v4.6" +_DEFAULT_RIFE_FILENAME = "flownet_v4.6.pkl" + + +def _load_rife_model(device, model_name=None): + """Load RIFE model and return a callable that interpolates two frames. + + Args: + device: Target device string ("cuda", "mps", "cpu"). + model_name: Optional HuggingFace repo ID containing RIFE weights. + Defaults to styler00dollar/RIFE-v4.6. + + Returns: + Callable that takes (frame1: PIL.Image, frame2: PIL.Image) -> PIL.Image + """ + from .rife_model import IFNet + from huggingface_hub import hf_hub_download + + repo_id = model_name if model_name is not None else _DEFAULT_RIFE_REPO + model_path = hf_hub_download(repo_id=repo_id, filename=_DEFAULT_RIFE_FILENAME) + + logger.info(f"Loading RIFE IFNet v4.6 to {device}") + + net = IFNet() + state_dict = torch.load(model_path, map_location="cpu", weights_only=True) + + # Strip "module." prefix that comes from DataParallel-saved checkpoints + cleaned = {} + for k, v in state_dict.items(): + cleaned[k.removeprefix("module.")] = v + + net.load_state_dict(cleaned) + net.to(device) + + def inference(img1, img2): + """Interpolate a single frame between two input frames.""" + arr1 = np.array(img1.convert("RGB")).astype(np.float32) / 255.0 + arr2 = np.array(img2.convert("RGB")).astype(np.float32) / 255.0 + t1 = torch.from_numpy(arr1).permute(2, 0, 1).unsqueeze(0).to(device) + t2 = torch.from_numpy(arr2).permute(2, 0, 1).unsqueeze(0).to(device) + + h, w = t1.shape[2], t1.shape[3] + ph = ((h - 1) // 32 + 1) * 32 + pw = ((w - 1) // 32 + 1) * 32 + padding = (0, pw - w, 0, ph - h) + t1_padded = torch.nn.functional.pad(t1, padding) + t2_padded = torch.nn.functional.pad(t2, padding) + + # Precompute warp grid and flow divisors for the padded resolution + tenFlow_div = torch.tensor([(pw - 1.0) / 2.0, (ph - 1.0) / 2.0], device=device) + backwarp_tenGrid = torch.cat( + [ + torch.linspace(-1.0, 1.0, pw, device=device) + .view(1, 1, 1, pw) + .expand(-1, -1, ph, -1), + torch.linspace(-1.0, 1.0, ph, device=device) + .view(1, 1, ph, 1) + .expand(-1, -1, -1, pw), + ], + 1, + ) + + timestep = torch.full((1, 1, ph, pw), 0.5, dtype=torch.float32, device=device) + scale_list = [8, 4, 2, 1] + + with torch.inference_mode(): + _, _, merged = net( + t1_padded, + t2_padded, + timestep, + scale_list, + tenFlow_div, + backwarp_tenGrid, + ) + + mid = merged[3][:, :, :h, :w] + mid = mid.squeeze(0).permute(1, 2, 0) + mid = (mid.clamp(0, 1) * 255).byte().cpu().numpy() + return Image.fromarray(mid) + + return inference diff --git a/dw/tasks/rife_model.py b/dw/tasks/rife_model.py new file mode 100644 index 00000000..d17bf679 --- /dev/null +++ b/dw/tasks/rife_model.py @@ -0,0 +1,178 @@ +""" +RIFE IFNet v4.6 — Real-Time Intermediate Flow Estimation. + +Vendored from https://github.com/hzwer/Practical-RIFE +Original architecture by Zhewei Huang et al. + +MIT License — Copyright (c) Megvii Inc. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def warp(tenInput, tenFlow, tenFlow_div, backwarp_tenGrid): + """Backward warp tenInput according to optical flow tenFlow.""" + tenFlow = torch.cat( + [ + tenFlow[:, 0:1] / tenFlow_div[0], + tenFlow[:, 1:2] / tenFlow_div[1], + ], + 1, + ) + g = (backwarp_tenGrid + tenFlow).permute(0, 2, 3, 1) + return F.grid_sample( + input=tenInput, + grid=g, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +class ResConv(nn.Module): + def __init__(self, c, dilation=1): + super().__init__() + self.conv = nn.Conv2d( + c, c, 3, 1, dilation, dilation=dilation, groups=1, bias=True + ) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x): + return self.relu(self.conv(x) * self.beta + x) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super().__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 2, 1), + conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1), + nn.PixelShuffle(2), + ) + + def forward(self, x, flow=None, scale=1): + x = F.interpolate( + x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) + if flow is not None: + flow = F.interpolate( + flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) * (1.0 / scale) + x = torch.cat((x, flow), 1) + feat = self.conv0(x) + feat = self.convblock(feat) + tmp = self.lastconv(feat) + tmp = F.interpolate( + tmp, scale_factor=scale, mode="bilinear", align_corners=False + ) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + return flow, mask + + +class IFNet(nn.Module): + def __init__(self): + super().__init__() + self.block0 = IFBlock(7 + 16, c=192) + self.block1 = IFBlock(8 + 4 + 16, c=128) + self.block2 = IFBlock(8 + 4 + 16, c=96) + self.block3 = IFBlock(8 + 4 + 16, c=64) + self.encode = nn.Sequential( + nn.Conv2d(3, 16, 3, 2, 1), + nn.ConvTranspose2d(16, 4, 4, 2, 1), + nn.LeakyReLU(0.2, True), + nn.Conv2d(4, 16, 3, 1, 1), + ) + + def forward(self, img0, img1, timestep, scale_list, tenFlow_div, backwarp_tenGrid): + f0 = self.encode(img0[:, :3]) + f1 = self.encode(img1[:, :3]) + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + + flow = None + for i, (block, scale) in enumerate( + zip( + [self.block0, self.block1, self.block2, self.block3], + scale_list, + ) + ): + if flow is None: + flow, mask = block( + torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), + None, + scale=scale, + ) + else: + wf0 = warp(f0, flow[:, :2], tenFlow_div, backwarp_tenGrid) + wf1 = warp(f1, flow[:, 2:4], tenFlow_div, backwarp_tenGrid) + fd, m = block( + torch.cat((warped_img0, warped_img1, wf0, wf1, timestep, mask), 1), + flow, + scale=scale, + ) + flow = flow + fd + mask = mask + m + + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0[:, :3], flow[:, :2], tenFlow_div, backwarp_tenGrid) + warped_img1 = warp(img1[:, :3], flow[:, 2:4], tenFlow_div, backwarp_tenGrid) + merged.append( + warped_img0 * torch.sigmoid(mask) + + warped_img1 * (1 - torch.sigmoid(mask)) + ) + + return flow_list, mask_list, merged diff --git a/dw/tasks/segment.py b/dw/tasks/segment.py new file mode 100644 index 00000000..cee27a1d --- /dev/null +++ b/dw/tasks/segment.py @@ -0,0 +1,108 @@ +""" +Image segmentation via GroundingDINO + SAM2. + +Takes an image and a text prompt, detects objects matching the prompt, +and returns a binary mask image (white = detected object). +""" + +import logging +import torch +import numpy as np +from PIL import Image, ImageOps +from transformers import ( + AutoProcessor, + AutoModelForZeroShotObjectDetection, + Sam2Model, + Sam2Processor, +) + +logger = logging.getLogger("dw") + +_DEFAULT_DINO_MODEL = "IDEA-Research/grounding-dino-base" +_DEFAULT_SAM_MODEL = "facebook/sam2-hiera-large" + + +def segment_image(image, prompt, device="cpu", **kwargs): + """Segment objects matching a text prompt, returning a binary mask. + + Args: + image: PIL Image to segment + prompt: Text description of object(s) to detect (e.g., "dog") + device: Target device ("cuda", "mps", "cpu") + **kwargs: + model_name: GroundingDINO model ID + sam_model_name: SAM2 model ID + threshold: Detection confidence threshold (default: 0.3) + invert: Invert the output mask (default: False) + + Returns: + PIL Image in mode "L" — white (255) for detected objects, black (0) for background. + """ + model_name = kwargs.get("model_name", _DEFAULT_DINO_MODEL) + sam_model_name = kwargs.get("sam_model_name", _DEFAULT_SAM_MODEL) + threshold = kwargs.get("threshold", 0.3) + invert = kwargs.get("invert", False) + + width, height = image.size + + logger.info(f"Loading GroundingDINO from {model_name}") + dino_processor = AutoProcessor.from_pretrained(model_name) + dino_model = AutoModelForZeroShotObjectDetection.from_pretrained(model_name).to( + device + ) + + inputs = dino_processor(images=image, text=prompt, return_tensors="pt").to(device) + + with torch.inference_mode(): + outputs = dino_model(**inputs) + + results = dino_processor.post_process_grounded_object_detection( + outputs, + inputs["input_ids"], + threshold=threshold, + target_sizes=[(height, width)], + ) + + boxes = results[0]["boxes"] + scores = results[0]["scores"] + labels = results[0]["labels"] + + logger.info(f"Detected {len(boxes)} objects: {labels} (scores: {scores.tolist()})") + + if len(boxes) == 0: + mask_image = Image.new("L", (width, height), 0) + if invert: + mask_image = ImageOps.invert(mask_image) + return mask_image + + logger.info(f"Loading SAM2 from {sam_model_name}") + sam_processor = Sam2Processor.from_pretrained(sam_model_name) + sam_model = Sam2Model.from_pretrained(sam_model_name).to(device) + + input_boxes = [boxes.cpu().tolist()] + + sam_inputs = sam_processor( + images=image, + input_boxes=input_boxes, + return_tensors="pt", + ).to(device) + + with torch.inference_mode(): + sam_outputs = sam_model(**sam_inputs) + + masks = sam_processor.post_process_masks( + sam_outputs.pred_masks, + sam_inputs["original_sizes"], + sam_inputs["reshaped_input_sizes"], + ) + + combined = masks[0][:, 0].sum(dim=0).clamp(0, 1) + mask_array = (combined.cpu().numpy() * 255).astype(np.uint8) + + mask_image = Image.fromarray(mask_array, mode="L") + + if invert: + mask_image = ImageOps.invert(mask_image) + + logger.info(f"Generated segmentation mask {width}x{height}") + return mask_image diff --git a/dw/tasks/task.py b/dw/tasks/task.py index 2777c87c..ef51eabe 100644 --- a/dw/tasks/task.py +++ b/dw/tasks/task.py @@ -11,6 +11,11 @@ ) from .upscale import upscale_image from .restore_faces import restore_faces +from .segment import segment_image +from .interpolate_frames import interpolate_frames +from .image_to_text import image_to_text +from .text_generation import generate_text +from .diffusion_upscale import diffusion_upscale logger = logging.getLogger("dw") @@ -90,6 +95,14 @@ def _handle_upscale(task, arguments, previous_pipelines): return upscale_image(image, model_name, device=task.device, **arguments) +@register_command("diffusion_upscale") +def _handle_diffusion_upscale(task, arguments, previous_pipelines): + """Upscale an image using a diffusion-based upscale pipeline""" + logger.debug("Diffusion upscaling image") + image = arguments.pop("image") + return diffusion_upscale(image, device=task.device, **arguments) + + @register_command("restore_faces") def _handle_restore_faces(task, arguments, previous_pipelines): """Restore faces in an image using a spandrel-compatible face restoration model""" @@ -99,6 +112,39 @@ def _handle_restore_faces(task, arguments, previous_pipelines): return restore_faces(image, model_name, device=task.device, **arguments) +@register_command("segment") +def _handle_segment(task, arguments, previous_pipelines): + """Segment objects in an image using text prompt""" + logger.debug("Segmenting image") + image = arguments.pop("image") + prompt = arguments.pop("prompt") + return segment_image(image, prompt, device=task.device, **arguments) + + +@register_command("interpolate_frames") +def _handle_interpolate_frames(task, arguments, previous_pipelines): + """Interpolate video frames to increase frame rate""" + logger.debug("Interpolating frames") + video = arguments.pop("video") + return interpolate_frames(video, device=task.device, **arguments) + + +@register_command("image_to_text") +def _handle_image_to_text(task, arguments, previous_pipelines): + """Generate text caption from an image""" + logger.debug("Captioning image") + image = arguments.pop("image") + return image_to_text(image, device=task.device, **arguments) + + +@register_command("text_generation") +def _handle_text_generation(task, arguments, previous_pipelines): + """Generate text from a prompt using a local LLM""" + logger.debug("Generating text") + prompt = arguments.pop("prompt") + return generate_text(prompt, device=task.device, **arguments) + + @register_command("batch_decode_post_process") def _handle_batch_decode(task, arguments, previous_pipelines): """Batch decode post-processing with pipeline reference""" diff --git a/dw/tasks/text_generation.py b/dw/tasks/text_generation.py new file mode 100644 index 00000000..ef9af3b4 --- /dev/null +++ b/dw/tasks/text_generation.py @@ -0,0 +1,60 @@ +""" +Text generation via HuggingFace transformers text-generation pipeline. + +Takes a prompt (and optional system prompt) and generates text using a +local language model. Useful for prompt expansion, rewriting, and +other text-to-text tasks. +""" + +import logging +import torch +from transformers import pipeline as hf_pipeline + +logger = logging.getLogger("dw") + +_DEFAULT_MODEL = "Qwen/Qwen2.5-1.5B-Instruct" + + +def generate_text(prompt, device="cpu", **kwargs): + """Generate text from a prompt using a local language model. + + Args: + prompt: The user message / prompt to expand or transform. + device: Target device ("cuda", "mps", "cpu"). + **kwargs: + model_name: HuggingFace model ID (default: Qwen/Qwen2.5-1.5B-Instruct). + system_prompt: Optional system instruction for the model. + max_new_tokens: Max tokens to generate (default: 500). + + Returns: + Generated text string. + """ + model_name = kwargs.get("model_name", _DEFAULT_MODEL) + system_prompt = kwargs.get("system_prompt", None) + max_new_tokens = int(kwargs.get("max_new_tokens", 500)) + + logger.info(f"Generating text with {model_name} on {device}") + + dtype = torch.float16 if device == "cuda" else torch.float32 + pipe = hf_pipeline( + "text-generation", + model=model_name, + device_map=device, + torch_dtype=dtype, + ) + + messages = [] + if system_prompt is not None: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + + results = pipe( + messages, + max_new_tokens=max_new_tokens, + return_full_text=False, + do_sample=False, + ) + + text = results[0]["generated_text"].strip() + logger.info(f"Generated: {text[:100]}{'...' if len(text) > 100 else ''}") + return text diff --git a/dw/workflow_schema.json b/dw/workflow_schema.json index 469e09f3..4c0d9175 100644 --- a/dw/workflow_schema.json +++ b/dw/workflow_schema.json @@ -672,6 +672,11 @@ "description": "Audio sample rate - only used when output is audio", "type": "integer", "default": 44100 + }, + "embed_metadata": { + "description": "Whether to embed generation parameters as metadata in saved images (PNG info chunks or EXIF). Only applies to image content types.", + "type": "boolean", + "default": false } }, "additionalProperties": { diff --git a/dw/workflows/augment_prompt.json b/dw/workflows/augment_prompt.json index ab3a805c..fb3f59f8 100644 --- a/dw/workflows/augment_prompt.json +++ b/dw/workflows/augment_prompt.json @@ -30,7 +30,7 @@ }, "from_pretrained_arguments": { "model_name": "microsoft/Phi-3.5-mini-instruct", - "device_map": "cuda", + "device_map": "auto", "torch_dtype": "{auto}", "trust_remote_code": true } diff --git a/dw/workflows/describe_image.json b/dw/workflows/describe_image.json index ee0b18e1..ed91705c 100644 --- a/dw/workflows/describe_image.json +++ b/dw/workflows/describe_image.json @@ -5,7 +5,7 @@ "id": "describe_image", "steps": [ { - "name": "desscribe_image_processor", + "name": "describe_image_processor", "pipeline": { "configuration": { "component_type": "transformers.AutoProcessor", @@ -35,8 +35,8 @@ "trust_remote_code": true }, "arguments": { - "input_ids": "previous_result:desscribe_image_processor.input_ids", - "pixel_values": "previous_result:desscribe_image_processor.pixel_values", + "input_ids": "previous_result:describe_image_processor.input_ids", + "pixel_values": "previous_result:describe_image_processor.pixel_values", "max_new_tokens": 4096, "num_beams": 3, "do_sample": false @@ -47,7 +47,7 @@ "name": "decode_image_description", "task": { "command": "batch_decode_post_process", - "pipeline_reference": "desscribe_image_processor", + "pipeline_reference": "describe_image_processor", "arguments": { "generated_ids": "previous_result:describe_image_model.generated_ids", "task": "" diff --git a/examples/CaptionToImage.json b/examples/CaptionToImage.json new file mode 100644 index 00000000..d7ef1ad4 --- /dev/null +++ b/examples/CaptionToImage.json @@ -0,0 +1,54 @@ +{ + "id": "CaptionToImage", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true" + }, + "steps": [ + { + "name": "input_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": [ + "variable:image_url" + ] + } + } + }, + { + "name": "caption", + "task": { + "command": "image_to_text", + "arguments": { + "image": "previous_result:input_image", + "max_new_tokens": 75 + } + }, + "result": { + "content_type": "text/plain", + "save": false + } + }, + { + "name": "generate", + "pipeline": { + "configuration": { + "offload": "model", + "component_type": "FluxPipeline" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "previous_result:caption", + "num_inference_steps": 25, + "guidance_scale": 3.5 + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} diff --git a/examples/DiffusionUpscaleX2.json b/examples/DiffusionUpscaleX2.json new file mode 100644 index 00000000..81526188 --- /dev/null +++ b/examples/DiffusionUpscaleX2.json @@ -0,0 +1,54 @@ +{ + "variables": { + "prompt": "a photo of an astronaut riding a horse on mars", + "upscale_prompt": "high quality, detailed", + "num_images_per_prompt": 1, + "num_inference_steps": 9, + "upscale_steps": 25, + "guidance_scale": 0.0, + "width": 512, + "height": 512 + }, + "id": "DiffusionUpscaleX2", + "steps": [ + { + "name": "generate", + "pipeline": { + "configuration": { + "component_type": "ZImagePipeline" + }, + "from_pretrained_arguments": { + "model_name": "Tongyi-MAI/Z-Image-Turbo", + "torch_dtype": "torch.bfloat16", + "low_cpu_mem_usage": false + }, + "arguments": { + "prompt": "variable:prompt", + "num_inference_steps": "variable:num_inference_steps", + "num_images_per_prompt": "variable:num_images_per_prompt", + "width": "variable:width", + "height": "variable:height", + "guidance_scale": "variable:guidance_scale" + } + }, + "result": { + "content_type": "image/jpeg" + } + }, + { + "name": "upscale", + "task": { + "command": "diffusion_upscale", + "arguments": { + "image": "previous_result:generate", + "prompt": "variable:upscale_prompt", + "mode": "x2", + "num_inference_steps": "variable:upscale_steps" + } + }, + "result": { + "content_type": "image/jpeg" + } + } + ] +} diff --git a/examples/DiffusionUpscaleX4.json b/examples/DiffusionUpscaleX4.json new file mode 100644 index 00000000..db53c055 --- /dev/null +++ b/examples/DiffusionUpscaleX4.json @@ -0,0 +1,56 @@ +{ + "variables": { + "prompt": "a photo of an astronaut riding a horse on mars", + "upscale_prompt": "high quality, detailed", + "negative_prompt": "blurry, low quality, artifacts", + "num_images_per_prompt": 1, + "num_inference_steps": 9, + "upscale_steps": 25, + "guidance_scale": 0.0, + "width": 512, + "height": 512 + }, + "id": "DiffusionUpscaleX4", + "steps": [ + { + "name": "generate", + "pipeline": { + "configuration": { + "component_type": "ZImagePipeline" + }, + "from_pretrained_arguments": { + "model_name": "Tongyi-MAI/Z-Image-Turbo", + "torch_dtype": "torch.bfloat16", + "low_cpu_mem_usage": false + }, + "arguments": { + "prompt": "variable:prompt", + "num_inference_steps": "variable:num_inference_steps", + "num_images_per_prompt": "variable:num_images_per_prompt", + "width": "variable:width", + "height": "variable:height", + "guidance_scale": "variable:guidance_scale" + } + }, + "result": { + "content_type": "image/jpeg" + } + }, + { + "name": "upscale", + "task": { + "command": "diffusion_upscale", + "arguments": { + "image": "previous_result:generate", + "prompt": "variable:upscale_prompt", + "negative_prompt": "variable:negative_prompt", + "mode": "x4", + "num_inference_steps": "variable:upscale_steps" + } + }, + "result": { + "content_type": "image/jpeg" + } + } + ] +} diff --git a/examples/ExpandAndGenerate.json b/examples/ExpandAndGenerate.json new file mode 100644 index 00000000..57cd0426 --- /dev/null +++ b/examples/ExpandAndGenerate.json @@ -0,0 +1,43 @@ +{ + "id": "ExpandAndGenerate", + "variables": { + "prompt": "a cat sitting on a windowsill" + }, + "steps": [ + { + "name": "expand", + "task": { + "command": "text_generation", + "arguments": { + "prompt": "variable:prompt", + "system_prompt": "You are a helpful AI assistant that creates detailed prompts for text to image generative AI. When supplied input generate only the prompt, no other text." + } + }, + "result": { + "content_type": "text/plain", + "save": false + } + }, + { + "name": "generate", + "pipeline": { + "configuration": { + "offload": "model", + "component_type": "FluxPipeline" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "previous_result:expand", + "num_inference_steps": 25, + "guidance_scale": 3.5 + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} diff --git a/examples/ExpandPrompt.json b/examples/ExpandPrompt.json new file mode 100644 index 00000000..a865d576 --- /dev/null +++ b/examples/ExpandPrompt.json @@ -0,0 +1,21 @@ +{ + "id": "ExpandPrompt", + "variables": { + "prompt": "a cat sitting on a windowsill" + }, + "steps": [ + { + "name": "expand", + "task": { + "command": "text_generation", + "arguments": { + "prompt": "variable:prompt", + "system_prompt": "You are a helpful AI assistant that creates detailed prompts for text to image generative AI. When supplied input generate only the prompt, no other text." + } + }, + "result": { + "content_type": "text/plain" + } + } + ] +} diff --git a/examples/ImageToText.json b/examples/ImageToText.json new file mode 100644 index 00000000..0a135ee4 --- /dev/null +++ b/examples/ImageToText.json @@ -0,0 +1,31 @@ +{ + "id": "ImageToText", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true" + }, + "steps": [ + { + "name": "input_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": [ + "variable:image_url" + ] + } + } + }, + { + "name": "caption", + "task": { + "command": "image_to_text", + "arguments": { + "image": "previous_result:input_image" + } + }, + "result": { + "content_type": "text/plain" + } + } + ] +} diff --git a/examples/ImageToTextBlip2.json b/examples/ImageToTextBlip2.json new file mode 100644 index 00000000..b4a65b85 --- /dev/null +++ b/examples/ImageToTextBlip2.json @@ -0,0 +1,35 @@ +{ + "id": "ImageToTextBlip2", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true", + "prompt": "Question: What is shown in this image? Answer:" + }, + "steps": [ + { + "name": "input_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": [ + "variable:image_url" + ] + } + } + }, + { + "name": "caption", + "task": { + "command": "image_to_text", + "arguments": { + "image": "previous_result:input_image", + "model_name": "Salesforce/blip2-opt-2.7b", + "prompt": "variable:prompt", + "max_new_tokens": 100 + } + }, + "result": { + "content_type": "text/plain" + } + } + ] +} diff --git a/examples/InterpolateFrames.json b/examples/InterpolateFrames.json new file mode 100644 index 00000000..71b69a84 --- /dev/null +++ b/examples/InterpolateFrames.json @@ -0,0 +1,50 @@ +{ + "id": "InterpolateFrames", + "variables": { + "prompt": "A cat walking across a sunlit room", + "multiplier": 2 + }, + "steps": [ + { + "name": "generate_video", + "pipeline": { + "configuration": { + "offload": "sequential", + "component_type": "MochiPipeline", + "vae": { + "enable_tiling": true, + "enable_slicing": true + } + }, + "from_pretrained_arguments": { + "model_name": "genmo/mochi-1-preview", + "variant": "bf16", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "variable:prompt", + "num_frames": 85 + } + }, + "result": { + "content_type": "video/mp4", + "save": false, + "fps": 30 + } + }, + { + "name": "interpolate", + "task": { + "command": "interpolate_frames", + "arguments": { + "video": "previous_result:generate_video", + "multiplier": "variable:multiplier" + } + }, + "result": { + "content_type": "video/mp4", + "fps": 60 + } + } + ] +} diff --git a/examples/MarigoldDepth.json b/examples/MarigoldDepth.json new file mode 100644 index 00000000..98b06105 --- /dev/null +++ b/examples/MarigoldDepth.json @@ -0,0 +1,40 @@ +{ + "id": "MarigoldDepth", + "variables": { + "image_url": "https://marigoldmonodepth.github.io/images/einstein.jpg" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "depth", + "pipeline": { + "configuration": { + "component_type": "MarigoldDepthPipeline" + }, + "from_pretrained_arguments": { + "model_name": "prs-eth/marigold-depth-lcm-v1-0", + "torch_dtype": "torch.float16", + "variant": "fp16" + }, + "arguments": { + "image": "previous_result:load_image" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} diff --git a/examples/MarigoldNormals.json b/examples/MarigoldNormals.json new file mode 100644 index 00000000..642096d5 --- /dev/null +++ b/examples/MarigoldNormals.json @@ -0,0 +1,40 @@ +{ + "id": "MarigoldNormals", + "variables": { + "image_url": "https://marigoldmonodepth.github.io/images/einstein.jpg" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "normals", + "pipeline": { + "configuration": { + "component_type": "MarigoldNormalsPipeline" + }, + "from_pretrained_arguments": { + "model_name": "prs-eth/marigold-normals-lcm-v1-0", + "torch_dtype": "torch.float16", + "variant": "fp16" + }, + "arguments": { + "image": "previous_result:load_image" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} diff --git a/examples/MetadataEmbed.json b/examples/MetadataEmbed.json new file mode 100644 index 00000000..3cc3a614 --- /dev/null +++ b/examples/MetadataEmbed.json @@ -0,0 +1,31 @@ +{ + "id": "MetadataEmbed", + "variables": { + "prompt": "A serene mountain landscape at golden hour, photorealistic", + "steps": 25 + }, + "steps": [ + { + "name": "generate", + "pipeline": { + "configuration": { + "offload": "model", + "component_type": "FluxPipeline" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "prompt": "variable:prompt", + "num_inference_steps": "variable:steps", + "guidance_scale": 3.5 + } + }, + "result": { + "content_type": "image/png", + "embed_metadata": true + } + } + ] +} diff --git a/examples/Segment.json b/examples/Segment.json new file mode 100644 index 00000000..19da0476 --- /dev/null +++ b/examples/Segment.json @@ -0,0 +1,35 @@ +{ + "id": "Segment", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/grounding_dino_example_input.png", + "prompt": "cat" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "segment", + "task": { + "command": "segment", + "arguments": { + "image": "previous_result:load_image", + "prompt": "variable:prompt" + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} diff --git a/examples/SegmentAndInpaint.json b/examples/SegmentAndInpaint.json new file mode 100644 index 00000000..a57afb57 --- /dev/null +++ b/examples/SegmentAndInpaint.json @@ -0,0 +1,63 @@ +{ + "id": "SegmentAndInpaint", + "variables": { + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/grounding_dino_example_input.png", + "segment_prompt": "cat", + "inpaint_prompt": "a golden retriever puppy sitting on the grass" + }, + "steps": [ + { + "name": "load_image", + "task": { + "command": "gather_images", + "arguments": { + "urls": ["variable:image_url"] + } + }, + "result": { + "content_type": "image/png", + "save": false + } + }, + { + "name": "segment", + "task": { + "command": "segment", + "arguments": { + "image": "previous_result:load_image", + "prompt": "variable:segment_prompt" + } + }, + "result": { + "content_type": "image/png", + "save": true + } + }, + { + "name": "inpaint", + "pipeline": { + "configuration": { + "component_type": "FluxFillPipeline", + "offload": "sequential" + }, + "from_pretrained_arguments": { + "model_name": "black-forest-labs/FLUX.1-Fill-dev", + "torch_dtype": "torch.bfloat16" + }, + "arguments": { + "image": "previous_result:load_image", + "mask_image": "previous_result:segment", + "prompt": "variable:inpaint_prompt", + "height": 1024, + "width": 1024, + "guidance_scale": 30, + "num_inference_steps": 50, + "max_sequence_length": 512 + } + }, + "result": { + "content_type": "image/png" + } + } + ] +} diff --git a/tests/test_diffusion_upscale.py b/tests/test_diffusion_upscale.py new file mode 100644 index 00000000..b7c3d830 --- /dev/null +++ b/tests/test_diffusion_upscale.py @@ -0,0 +1,188 @@ +"""Tests for diffusion_upscale task.""" + +import unittest +from unittest.mock import patch, MagicMock +from PIL import Image + +from dw.tasks.diffusion_upscale import diffusion_upscale, _MODELS + + +class TestDiffusionUpscale(unittest.TestCase): + """Tests for the diffusion_upscale function.""" + + def _make_image(self): + return Image.new("RGB", (128, 128), color="red") + + def _make_mock_pipeline(self): + """Create a mock pipeline that returns a PIL image.""" + output_image = Image.new("RGB", (512, 512), color="blue") + mock_result = MagicMock() + mock_result.images = [output_image] + + mock_pipe = MagicMock() + mock_pipe.return_value = mock_result + return mock_pipe + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_returns_pil_image(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + result = diffusion_upscale(self._make_image(), device="cpu") + + self.assertIsInstance(result, Image.Image) + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_default_mode_is_x4(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu") + + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.assert_called_once() + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_x2_mode_uses_latent_pipeline(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionLatentUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu", mode="x2") + + mock_diffusers.StableDiffusionLatentUpscalePipeline.from_pretrained.assert_called_once() + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_x4_includes_noise_level(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu") + + call_kwargs = mock_pipe.call_args[1] + self.assertIn("noise_level", call_kwargs) + self.assertEqual(call_kwargs["noise_level"], 20) + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_x2_excludes_noise_level(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionLatentUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu", mode="x2") + + call_kwargs = mock_pipe.call_args[1] + self.assertNotIn("noise_level", call_kwargs) + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_custom_prompt(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu", prompt="a photo of a cat") + + call_kwargs = mock_pipe.call_args[1] + self.assertEqual(call_kwargs["prompt"], "a photo of a cat") + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_negative_prompt(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale( + self._make_image(), device="cpu", negative_prompt="blurry, low quality" + ) + + call_kwargs = mock_pipe.call_args[1] + self.assertEqual(call_kwargs["negative_prompt"], "blurry, low quality") + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_no_negative_prompt_by_default(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu") + + call_kwargs = mock_pipe.call_args[1] + self.assertNotIn("negative_prompt", call_kwargs) + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_custom_inference_steps(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu", num_inference_steps=50) + + call_kwargs = mock_pipe.call_args[1] + self.assertEqual(call_kwargs["num_inference_steps"], 50) + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_custom_guidance_scale(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale(self._make_image(), device="cpu", guidance_scale=7.5) + + call_kwargs = mock_pipe.call_args[1] + self.assertEqual(call_kwargs["guidance_scale"], 7.5) + + @patch("dw.tasks.diffusion_upscale.diffusers") + def test_custom_model_name(self, mock_diffusers): + mock_pipe = self._make_mock_pipeline() + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.return_value = ( + mock_pipe + ) + + diffusion_upscale( + self._make_image(), device="cpu", model_name="my-org/my-upscaler" + ) + + call_args = ( + mock_diffusers.StableDiffusionUpscalePipeline.from_pretrained.call_args + ) + self.assertEqual(call_args[0][0], "my-org/my-upscaler") + + def test_invalid_mode_raises(self): + with self.assertRaises(ValueError) as ctx: + diffusion_upscale(self._make_image(), device="cpu", mode="x8") + self.assertIn("x8", str(ctx.exception)) + + def test_models_config_has_expected_modes(self): + self.assertIn("x4", _MODELS) + self.assertIn("x2", _MODELS) + self.assertEqual( + _MODELS["x4"]["pipeline_class"], "StableDiffusionUpscalePipeline" + ) + self.assertEqual( + _MODELS["x2"]["pipeline_class"], "StableDiffusionLatentUpscalePipeline" + ) + + +class TestDiffusionUpscaleRegistration(unittest.TestCase): + """Test that diffusion_upscale is registered as a task command.""" + + def test_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + self.assertIn("diffusion_upscale", _COMMAND_REGISTRY) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_to_text.py b/tests/test_image_to_text.py new file mode 100644 index 00000000..ffae144b --- /dev/null +++ b/tests/test_image_to_text.py @@ -0,0 +1,99 @@ +"""Tests for image_to_text task.""" + +import unittest +from unittest.mock import patch, MagicMock +from PIL import Image + +from dw.tasks.image_to_text import image_to_text, _DEFAULT_MODEL + + +class TestImageToText(unittest.TestCase): + """Tests for the image_to_text function.""" + + def _make_image(self): + return Image.new("RGB", (64, 64), color="red") + + @patch("dw.tasks.image_to_text.hf_pipeline") + def test_returns_caption_string(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "a red square"}] + mock_pipeline.return_value = pipe + + result = image_to_text(self._make_image(), device="cpu") + + self.assertEqual(result, "a red square") + mock_pipeline.assert_called_once() + + @patch("dw.tasks.image_to_text.hf_pipeline") + def test_uses_default_model(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "caption"}] + mock_pipeline.return_value = pipe + + image_to_text(self._make_image(), device="cpu") + + call_kwargs = mock_pipeline.call_args + self.assertEqual(call_kwargs[1]["model"], _DEFAULT_MODEL) + + @patch("dw.tasks.image_to_text.hf_pipeline") + def test_custom_model_name(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "detailed caption"}] + mock_pipeline.return_value = pipe + + image_to_text( + self._make_image(), + device="cpu", + model_name="Salesforce/blip2-opt-2.7b", + ) + + call_kwargs = mock_pipeline.call_args + self.assertEqual(call_kwargs[1]["model"], "Salesforce/blip2-opt-2.7b") + + @patch("dw.tasks.image_to_text.hf_pipeline") + def test_prompt_passed_as_generate_kwarg(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "a photo of a dog"}] + mock_pipeline.return_value = pipe + + image_to_text( + self._make_image(), + device="cpu", + prompt="Question: what is this? Answer:", + ) + + call_args = pipe.call_args + self.assertIn("prompt", call_args[1]["generate_kwargs"]) + + @patch("dw.tasks.image_to_text.hf_pipeline") + def test_max_new_tokens(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "caption"}] + mock_pipeline.return_value = pipe + + image_to_text(self._make_image(), device="cpu", max_new_tokens=100) + + call_args = pipe.call_args + self.assertEqual(call_args[1]["generate_kwargs"]["max_new_tokens"], 100) + + @patch("dw.tasks.image_to_text.hf_pipeline") + def test_strips_whitespace(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": " a caption with spaces "}] + mock_pipeline.return_value = pipe + + result = image_to_text(self._make_image(), device="cpu") + self.assertEqual(result, "a caption with spaces") + + +class TestImageToTextRegistration(unittest.TestCase): + """Test that image_to_text is registered as a task command.""" + + def test_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + self.assertIn("image_to_text", _COMMAND_REGISTRY) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_interpolate_frames.py b/tests/test_interpolate_frames.py new file mode 100644 index 00000000..039299dc --- /dev/null +++ b/tests/test_interpolate_frames.py @@ -0,0 +1,84 @@ +import pytest +from unittest.mock import patch, MagicMock +from PIL import Image +import numpy as np + + +def _make_test_frames(count=4, width=64, height=64): + """Create a list of test frames with different colors.""" + frames = [] + for i in range(count): + shade = int(255 * i / max(count - 1, 1)) + frames.append(Image.new("RGB", (width, height), color=(shade, shade, shade))) + return frames + + +class TestInterpolateFrames: + """Test interpolate_frames function.""" + + @patch("dw.tasks.interpolate_frames._load_rife_model") + def test_2x_doubles_frame_count(self, mock_load): + """2x multiplier should produce 2N-1 frames from N input frames.""" + from dw.tasks.interpolate_frames import interpolate_frames + + mock_model = MagicMock() + + def fake_inference(img1, img2): + arr1 = np.array(img1).astype(np.float32) + arr2 = np.array(img2).astype(np.float32) + mid = ((arr1 + arr2) / 2).astype(np.uint8) + return Image.fromarray(mid) + + mock_model.side_effect = fake_inference + mock_load.return_value = mock_model + + frames = _make_test_frames(4) + result = interpolate_frames(frames, multiplier=2) + + # 4 frames with 2x: (4-1)*2 + 1 = 7 + assert len(result) == 7 + assert all(isinstance(f, Image.Image) for f in result) + + @patch("dw.tasks.interpolate_frames._load_rife_model") + def test_4x_quadruples_frame_count(self, mock_load): + """4x multiplier should run two passes of 2x.""" + from dw.tasks.interpolate_frames import interpolate_frames + + mock_model = MagicMock() + + def fake_inference(img1, img2): + arr1 = np.array(img1).astype(np.float32) + arr2 = np.array(img2).astype(np.float32) + mid = ((arr1 + arr2) / 2).astype(np.uint8) + return Image.fromarray(mid) + + mock_model.side_effect = fake_inference + mock_load.return_value = mock_model + + frames = _make_test_frames(4) + result = interpolate_frames(frames, multiplier=4) + + # Two passes of 2x: 4 -> 7 -> 13 + assert len(result) == 13 + assert all(isinstance(f, Image.Image) for f in result) + + def test_invalid_multiplier_raises(self): + """multiplier must be 2, 4, or 8.""" + from dw.tasks.interpolate_frames import interpolate_frames + + with pytest.raises(ValueError, match="multiplier"): + interpolate_frames(_make_test_frames(2), multiplier=3) + + def test_single_frame_raises(self): + """Need at least 2 frames to interpolate.""" + from dw.tasks.interpolate_frames import interpolate_frames + + with pytest.raises(ValueError, match="at least 2"): + interpolate_frames(_make_test_frames(1), multiplier=2) + + +class TestInterpolateFramesRegistration: + def test_interpolate_frames_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + assert "interpolate_frames" in _COMMAND_REGISTRY diff --git a/tests/test_resize_bucket.py b/tests/test_resize_bucket.py new file mode 100644 index 00000000..d6128595 --- /dev/null +++ b/tests/test_resize_bucket.py @@ -0,0 +1,98 @@ +"""Tests for resize_bucket image processing command.""" + +import unittest +from PIL import Image + +from dw.tasks.image_utils import resize_bucket, _DEFAULT_RATIOS + + +class TestResizeBucket(unittest.TestCase): + """Tests for the resize_bucket function.""" + + def test_square_image_picks_1_1(self): + img = Image.new("RGB", (500, 500)) + result = resize_bucket(img, resolution=1024) + self.assertEqual(result.width, result.height) + + def test_landscape_picks_landscape_ratio(self): + # 1600x900 is 16:9 + img = Image.new("RGB", (1600, 900)) + result = resize_bucket(img, resolution=1024) + self.assertGreater(result.width, result.height) + + def test_portrait_picks_portrait_ratio(self): + # 900x1600 is 9:16 + img = Image.new("RGB", (900, 1600)) + result = resize_bucket(img, resolution=1024) + self.assertGreater(result.height, result.width) + + def test_dimensions_aligned_to_64(self): + img = Image.new("RGB", (1600, 900)) + result = resize_bucket(img, resolution=1024) + self.assertEqual(result.width % 64, 0) + self.assertEqual(result.height % 64, 0) + + def test_custom_alignment(self): + img = Image.new("RGB", (800, 600)) + result = resize_bucket(img, resolution=512, alignment=32) + self.assertEqual(result.width % 32, 0) + self.assertEqual(result.height % 32, 0) + + def test_custom_ratios(self): + # 1800x900 is 2:1, should pick [2, 1] over [1, 1] + img = Image.new("RGB", (1800, 900)) + result = resize_bucket(img, resolution=1024, ratios=[[1, 1], [2, 1]]) + self.assertGreater(result.width, result.height) + + def test_4_3_image(self): + # 800x600 is exactly 4:3 + img = Image.new("RGB", (800, 600)) + result = resize_bucket(img, resolution=1024) + ratio = result.width / result.height + # Should pick 4:3 (1.333) — verify it's close + self.assertAlmostEqual(ratio, 4 / 3, delta=0.1) + + def test_3_2_image(self): + # 1200x800 is 3:2 + img = Image.new("RGB", (1200, 800)) + result = resize_bucket(img, resolution=1024) + ratio = result.width / result.height + self.assertAlmostEqual(ratio, 3 / 2, delta=0.1) + + def test_resolution_controls_short_side(self): + img = Image.new("RGB", (500, 500)) + result = resize_bucket(img, resolution=512) + # 1:1 ratio, so both sides should be ~512 + self.assertEqual(result.width, 512) + self.assertEqual(result.height, 512) + + def test_converts_to_rgb(self): + img = Image.new("RGBA", (500, 500)) + result = resize_bucket(img, resolution=512) + self.assertEqual(result.mode, "RGB") + + def test_default_ratios_has_expected_entries(self): + # Sanity check that we have the standard ratios + ratio_values = {(r[0], r[1]) for r in _DEFAULT_RATIOS} + self.assertIn((1, 1), ratio_values) + self.assertIn((16, 9), ratio_values) + self.assertIn((9, 16), ratio_values) + self.assertIn((4, 3), ratio_values) + self.assertIn((3, 4), ratio_values) + + +class TestResizeBucketRegistration(unittest.TestCase): + """Test that resize_bucket is accessible via process_image.""" + + def test_process_image_dispatches(self): + from dw.tasks.image_utils import process_image + + img = Image.new("RGB", (800, 600)) + result = process_image(img, "resize_bucket", "cpu", {"resolution": 512}) + self.assertIsInstance(result, Image.Image) + self.assertEqual(result.width % 64, 0) + self.assertEqual(result.height % 64, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_result.py b/tests/test_result.py index f3d65981..5ba47387 100644 --- a/tests/test_result.py +++ b/tests/test_result.py @@ -190,5 +190,82 @@ def test_empty_content_type(self): assert guess_extension("") == "" +class TestMetadataEmbedding: + """Test opt-in metadata embedding in saved images.""" + + def test_png_metadata_embedded(self): + """When embed_metadata is true, PNG should contain parameters text chunk.""" + with tempfile.TemporaryDirectory() as temp_dir: + result_def = { + "content_type": "image/png", + "save": True, + "embed_metadata": True, + } + result = Result(result_def) + result.set_metadata( + { + "workflow_id": "test_workflow", + "step_name": "generate", + "model_name": "test/model", + "arguments": {"prompt": "a cat", "num_inference_steps": 25}, + } + ) + + img = Image.new("RGB", (64, 64), color=(128, 64, 32)) + result.add_result(img) + result.save(temp_dir, "test_output") + + output_file = os.path.join(temp_dir, "test_output-0.0.png") + assert os.path.exists(output_file) + + saved_img = Image.open(output_file) + assert "parameters" in saved_img.info + metadata = json.loads(saved_img.info["parameters"]) + assert metadata["workflow_id"] == "test_workflow" + assert metadata["arguments"]["prompt"] == "a cat" + + def test_no_metadata_when_not_enabled(self): + """When embed_metadata is absent/false, no metadata should be embedded.""" + with tempfile.TemporaryDirectory() as temp_dir: + result_def = {"content_type": "image/png", "save": True} + result = Result(result_def) + + img = Image.new("RGB", (64, 64), color=(128, 64, 32)) + result.add_result(img) + result.save(temp_dir, "test_output") + + output_file = os.path.join(temp_dir, "test_output-0.0.png") + saved_img = Image.open(output_file) + assert "parameters" not in saved_img.info + + def test_set_metadata_method(self): + """set_metadata should store metadata on the Result instance.""" + result = Result({}) + assert result.metadata is None + + metadata = {"workflow_id": "test", "step_name": "step1"} + result.set_metadata(metadata) + assert result.metadata == metadata + + def test_metadata_with_embed_false(self): + """When embed_metadata is explicitly false, no metadata embedded even if set.""" + with tempfile.TemporaryDirectory() as temp_dir: + result_def = { + "content_type": "image/png", + "save": True, + "embed_metadata": False, + } + result = Result(result_def) + result.set_metadata({"workflow_id": "test"}) + + img = Image.new("RGB", (64, 64), color=(128, 64, 32)) + result.add_result(img) + result.save(temp_dir, "test_output") + + output_file = os.path.join(temp_dir, "test_output-0.0.png") + saved_img = Image.open(output_file) + assert "parameters" not in saved_img.info + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_segment.py b/tests/test_segment.py new file mode 100644 index 00000000..fb1bba5c --- /dev/null +++ b/tests/test_segment.py @@ -0,0 +1,156 @@ +import pytest +from unittest.mock import patch, MagicMock +import torch +import numpy as np +from PIL import Image + + +def _make_test_image(width=640, height=480): + """Create a simple test image.""" + return Image.new("RGB", (width, height), color=(128, 64, 32)) + + +def _make_batch_encoding(data): + """Create a MagicMock that behaves like a BatchEncoding (dict-like + .to()).""" + mock = MagicMock() + mock.__getitem__ = lambda self, key: data[key] + mock.__contains__ = lambda self, key: key in data + mock.to.return_value = mock + return mock + + +class TestSegmentImage: + """Test segment_image function with mocked models.""" + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_returns_pil_image_mode_l( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + """segment_image should return a grayscale PIL Image.""" + from dw.tasks.segment import segment_image + + mock_processor_instance = MagicMock() + mock_auto_proc.from_pretrained.return_value = mock_processor_instance + dino_inputs = _make_batch_encoding({"input_ids": torch.zeros(1, 10)}) + mock_processor_instance.return_value = dino_inputs + mock_processor_instance.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.tensor([[100.0, 100.0, 300.0, 300.0]]), + "scores": torch.tensor([0.9]), + "labels": ["dog"], + } + ] + + mock_model_instance = MagicMock() + mock_auto_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.to.return_value = mock_model_instance + mock_model_instance.return_value = MagicMock() + + mock_sam_proc_instance = MagicMock() + mock_sam_proc.from_pretrained.return_value = mock_sam_proc_instance + sam_inputs = _make_batch_encoding( + { + "pixel_values": torch.zeros(1, 3, 256, 256), + "input_boxes": torch.tensor([[[100.0, 100.0, 300.0, 300.0]]]), + "original_sizes": torch.tensor([[480, 640]]), + "reshaped_input_sizes": torch.tensor([[256, 256]]), + } + ) + mock_sam_proc_instance.return_value = sam_inputs + + mock_sam_model_instance = MagicMock() + mock_sam_model.from_pretrained.return_value = mock_sam_model_instance + mock_sam_model_instance.to.return_value = mock_sam_model_instance + mask_tensor = torch.zeros(1, 1, 3, 480, 640) + mask_tensor[0, 0, 0, 100:300, 100:300] = 1.0 + mock_sam_model_instance.return_value = MagicMock(pred_masks=mask_tensor) + + post_mask = torch.zeros(1, 1, 480, 640) + post_mask[0, 0, 100:300, 100:300] = 1.0 + mock_sam_proc_instance.post_process_masks.return_value = [post_mask] + + image = _make_test_image() + result = segment_image(image, "dog") + + assert isinstance(result, Image.Image) + assert result.mode == "L" + assert result.size == (640, 480) + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_no_detections_returns_black_mask( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + """When nothing is detected, return an all-black mask.""" + from dw.tasks.segment import segment_image + + mock_processor_instance = MagicMock() + mock_auto_proc.from_pretrained.return_value = mock_processor_instance + dino_inputs = _make_batch_encoding({"input_ids": torch.zeros(1, 10)}) + mock_processor_instance.return_value = dino_inputs + mock_processor_instance.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.zeros(0, 4), + "scores": torch.zeros(0), + "labels": [], + } + ] + + mock_model_instance = MagicMock() + mock_auto_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.to.return_value = mock_model_instance + mock_model_instance.return_value = MagicMock() + + image = _make_test_image() + result = segment_image(image, "nonexistent_object") + + assert isinstance(result, Image.Image) + assert result.mode == "L" + arr = np.array(result) + assert arr.max() == 0 + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_invert_flag( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + """When invert=True, mask should be inverted.""" + from dw.tasks.segment import segment_image + + mock_processor_instance = MagicMock() + mock_auto_proc.from_pretrained.return_value = mock_processor_instance + dino_inputs = _make_batch_encoding({"input_ids": torch.zeros(1, 10)}) + mock_processor_instance.return_value = dino_inputs + mock_processor_instance.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.zeros(0, 4), + "scores": torch.zeros(0), + "labels": [], + } + ] + + mock_model_instance = MagicMock() + mock_auto_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.to.return_value = mock_model_instance + mock_model_instance.return_value = MagicMock() + + image = _make_test_image() + result = segment_image(image, "nonexistent_object", invert=True) + + assert isinstance(result, Image.Image) + arr = np.array(result) + assert arr.min() == 255 + + +class TestSegmentTaskRegistration: + def test_segment_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + assert "segment" in _COMMAND_REGISTRY diff --git a/tests/test_step.py b/tests/test_step.py index 208a520f..620b4b38 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -154,6 +154,147 @@ def capture_args(args, pipelines): assert ("img2.jpg", "prompt A") in combinations assert ("img2.jpg", "prompt B") in combinations + # --- embed_metadata tests --- + + def test_embed_metadata_disabled_by_default(self): + """embed_metadata not set → Result.metadata stays None""" + step_def = {"name": "test_step", "result": {}} + step = Step(step_def, default_seed=42) + + mock_action = Mock() + mock_action.name = "mock_action" + mock_action.argument_template = {"prompt": "test"} + mock_action.run = Mock(return_value="result_value") + + result = step.run({}, {}, mock_action) + + assert result.metadata is None + + def test_embed_metadata_false_leaves_metadata_none(self): + """embed_metadata explicitly False → Result.metadata stays None""" + step_def = {"name": "test_step", "result": {"embed_metadata": False}} + step = Step(step_def, default_seed=42) + + mock_action = Mock() + mock_action.name = "mock_action" + mock_action.argument_template = {"prompt": "test"} + mock_action.run = Mock(return_value="result_value") + + result = step.run({}, {}, mock_action) + + assert result.metadata is None + + def test_embed_metadata_true_pipeline_step(self): + """embed_metadata=True for a pipeline step → Result carries expected metadata""" + step_def = { + "name": "gen_step", + "pipeline": { + "from_pretrained_arguments": {"model_name": "my-org/my-model"}, + "arguments": {"prompt": "a cat", "num_inference_steps": 25}, + }, + "result": {"embed_metadata": True}, + } + step = Step(step_def, default_seed=42) + + mock_action = Mock() + mock_action.name = "mock_action" + mock_action.argument_template = {"prompt": "a cat"} + mock_action.run = Mock(return_value="img.png") + + result = step.run({}, {}, mock_action) + + assert result.metadata is not None + assert result.metadata["step_name"] == "gen_step" + assert result.metadata["model_name"] == "my-org/my-model" + assert result.metadata["arguments"] == {"prompt": "a cat", "num_inference_steps": 25} + + def test_embed_metadata_true_task_step(self): + """embed_metadata=True for a task step → Result carries task metadata""" + step_def = { + "name": "proc_step", + "task": { + "command": "process_image", + "arguments": {"operation": "resize", "width": 512}, + }, + "result": {"embed_metadata": True}, + } + step = Step(step_def, default_seed=42) + + mock_action = Mock() + mock_action.name = "mock_action" + mock_action.argument_template = {"operation": "resize"} + mock_action.run = Mock(return_value="out.png") + + result = step.run({}, {}, mock_action) + + assert result.metadata is not None + assert result.metadata["step_name"] == "proc_step" + assert result.metadata["task_command"] == "process_image" + assert result.metadata["arguments"] == {"operation": "resize", "width": 512} + + def test_embed_metadata_pipeline_without_model_name(self): + """embed_metadata=True for pipeline step with no model_name → no model_name key""" + step_def = { + "name": "anon_step", + "pipeline": { + "from_pretrained_arguments": {}, + "arguments": {"prompt": "test"}, + }, + "result": {"embed_metadata": True}, + } + step = Step(step_def, default_seed=42) + + mock_action = Mock() + mock_action.name = "mock_action" + mock_action.argument_template = {"prompt": "test"} + mock_action.run = Mock(return_value="out.png") + + result = step.run({}, {}, mock_action) + + assert result.metadata is not None + assert "model_name" not in result.metadata + assert result.metadata["step_name"] == "anon_step" + + def test_embed_metadata_shared_across_all_iterations(self): + """With embed_metadata=True and multiple iterations the single Result carries metadata""" + step_def = { + "name": "multi_step", + "pipeline": { + "from_pretrained_arguments": {"model_name": "org/model"}, + "arguments": {"image": "previous_result:images"}, + }, + "result": {"embed_metadata": True}, + } + step = Step(step_def, default_seed=42) + + mock_action = Mock() + mock_action.name = "mock_action" + mock_action.argument_template = {"image": "previous_result:images"} + + call_count = 0 + + def mock_run(args, pipelines): + nonlocal call_count + call_count += 1 + return f"result_{call_count}" + + mock_action.run = Mock(side_effect=mock_run) + + images_result = Result({}) + images_result.add_result(["img1.jpg", "img2.jpg", "img3.jpg"]) + previous_results = {"images": images_result} + + result = step.run(previous_results, {}, mock_action) + + # Three iterations should have all run + assert mock_action.run.call_count == 3 + assert len(result.result_list) == 3 + + # Metadata is set once on the Result (not per-iteration) + assert result.metadata is not None + assert result.metadata["step_name"] == "multi_step" + assert result.metadata["model_name"] == "org/model" + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_strip_exif_and_watermark.py b/tests/test_strip_exif_and_watermark.py new file mode 100644 index 00000000..ec56964e --- /dev/null +++ b/tests/test_strip_exif_and_watermark.py @@ -0,0 +1,108 @@ +"""Tests for strip_exif and add_watermark image processing commands.""" + +import unittest +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +from dw.tasks.image_utils import strip_exif, add_watermark, process_image + + +class TestStripExif(unittest.TestCase): + """Tests for the strip_exif function.""" + + def test_returns_image_same_size(self): + img = Image.new("RGB", (200, 100), color="red") + result = strip_exif(img) + self.assertEqual(result.size, (200, 100)) + + def test_preserves_pixel_data(self): + img = Image.new("RGB", (2, 2), color=(255, 0, 0)) + result = strip_exif(img) + # Verify pixel content is preserved + self.assertEqual(result.getpixel((0, 0)), (255, 0, 0)) + self.assertEqual(result.getpixel((1, 1)), (255, 0, 0)) + + def test_removes_png_metadata(self): + img = Image.new("RGB", (10, 10)) + info = PngInfo() + info.add_text("Comment", "secret location data") + # Simulate an image with metadata by setting .info + img.info["Comment"] = "secret location data" + + result = strip_exif(img) + # The clean image should have no info dict entries carried over + self.assertNotIn("Comment", result.info) + + def test_preserves_mode(self): + img = Image.new("RGBA", (10, 10), color=(255, 0, 0, 128)) + result = strip_exif(img) + self.assertEqual(result.mode, "RGBA") + + def test_dispatch_via_process_image(self): + img = Image.new("RGB", (50, 50)) + result = process_image(img, "strip_exif", "cpu", {}) + self.assertIsInstance(result, Image.Image) + self.assertEqual(result.size, (50, 50)) + + +class TestAddWatermark(unittest.TestCase): + """Tests for the add_watermark function.""" + + def test_returns_rgb_image(self): + img = Image.new("RGB", (200, 100), color="blue") + result = add_watermark(img) + self.assertEqual(result.mode, "RGB") + self.assertEqual(result.size, (200, 100)) + + def test_modifies_pixels(self): + img = Image.new("RGB", (200, 100), color=(0, 0, 0)) + result = add_watermark(img, text="WATERMARK", opacity=255) + # The watermarked image should differ from the all-black original + self.assertNotEqual(img.tobytes(), result.tobytes()) + + def test_default_text(self): + # Should not raise with defaults + img = Image.new("RGB", (400, 200)) + result = add_watermark(img) + self.assertIsInstance(result, Image.Image) + + def test_custom_text(self): + img = Image.new("RGB", (400, 200)) + result = add_watermark(img, text="DO NOT DISTRIBUTE") + self.assertIsInstance(result, Image.Image) + + def test_all_positions(self): + img = Image.new("RGB", (400, 200)) + for pos in ["bottom-right", "bottom-left", "top-right", "top-left", "center"]: + result = add_watermark(img, position=pos) + self.assertEqual(result.size, (400, 200)) + + def test_invalid_position_falls_back(self): + img = Image.new("RGB", (400, 200)) + # Unknown position should fall back to bottom-right + result = add_watermark(img, position="nonsense") + self.assertIsInstance(result, Image.Image) + + def test_custom_color(self): + img = Image.new("RGB", (400, 200)) + result = add_watermark(img, color=(255, 0, 0)) + self.assertIsInstance(result, Image.Image) + + def test_custom_font_size(self): + img = Image.new("RGB", (400, 200)) + result = add_watermark(img, font_size=24) + self.assertIsInstance(result, Image.Image) + + def test_rgba_input_converted(self): + img = Image.new("RGBA", (200, 100)) + result = add_watermark(img) + self.assertEqual(result.mode, "RGB") + + def test_dispatch_via_process_image(self): + img = Image.new("RGB", (200, 100)) + result = process_image(img, "add_watermark", "cpu", {"text": "TEST"}) + self.assertIsInstance(result, Image.Image) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py new file mode 100644 index 00000000..3789381f --- /dev/null +++ b/tests/test_text_generation.py @@ -0,0 +1,109 @@ +"""Tests for text_generation task.""" + +import unittest +from unittest.mock import patch, MagicMock + +from dw.tasks.text_generation import generate_text, _DEFAULT_MODEL + + +class TestTextGeneration(unittest.TestCase): + """Tests for the generate_text function.""" + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_returns_generated_string(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "an expanded detailed prompt"}] + mock_pipeline.return_value = pipe + + result = generate_text("a cat", device="cpu") + + self.assertEqual(result, "an expanded detailed prompt") + mock_pipeline.assert_called_once() + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_uses_default_model(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "output"}] + mock_pipeline.return_value = pipe + + generate_text("test", device="cpu") + + call_kwargs = mock_pipeline.call_args + self.assertEqual(call_kwargs[1]["model"], _DEFAULT_MODEL) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_custom_model_name(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "output"}] + mock_pipeline.return_value = pipe + + generate_text("test", device="cpu", model_name="meta-llama/Llama-3.2-1B-Instruct") + + call_kwargs = mock_pipeline.call_args + self.assertEqual(call_kwargs[1]["model"], "meta-llama/Llama-3.2-1B-Instruct") + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_system_prompt_included(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "output"}] + mock_pipeline.return_value = pipe + + generate_text( + "a cat", + device="cpu", + system_prompt="You expand prompts for image generation.", + ) + + call_args = pipe.call_args + messages = call_args[0][0] + self.assertEqual(len(messages), 2) + self.assertEqual(messages[0]["role"], "system") + self.assertEqual(messages[0]["content"], "You expand prompts for image generation.") + self.assertEqual(messages[1]["role"], "user") + self.assertEqual(messages[1]["content"], "a cat") + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_no_system_prompt(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "output"}] + mock_pipeline.return_value = pipe + + generate_text("a cat", device="cpu") + + call_args = pipe.call_args + messages = call_args[0][0] + self.assertEqual(len(messages), 1) + self.assertEqual(messages[0]["role"], "user") + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_max_new_tokens(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": "output"}] + mock_pipeline.return_value = pipe + + generate_text("test", device="cpu", max_new_tokens=200) + + call_args = pipe.call_args + self.assertEqual(call_args[1]["max_new_tokens"], 200) + + @patch("dw.tasks.text_generation.hf_pipeline") + def test_strips_whitespace(self, mock_pipeline): + pipe = MagicMock() + pipe.return_value = [{"generated_text": " some text with spaces \n"}] + mock_pipeline.return_value = pipe + + result = generate_text("test", device="cpu") + self.assertEqual(result, "some text with spaces") + + +class TestTextGenerationRegistration(unittest.TestCase): + """Test that text_generation is registered as a task command.""" + + def test_command_registered(self): + from dw.tasks.task import _COMMAND_REGISTRY + + self.assertIn("text_generation", _COMMAND_REGISTRY) + + +if __name__ == "__main__": + unittest.main()