Codestin Search App

Carlofkl · 2026-05-27T06:16:54Z

Context

This PR integrates DreamLite — ByteDance's text-to-image / image-edit diffusion model — into diffusers, following an invitation from @NielsRogge to release the model on the Hub in diffusers format.

Related issue: ByteVisionLab/DreamLite#3 (comment)

Model cards (public, ungated):

Base (3-branch dual CFG): https://huggingface.co/carlofkl/DreamLite-base
Mobile (distilled, single forward): https://huggingface.co/carlofkl/DreamLite-mobile

Both repos use a diffusers branch (loaded via revision="diffusers") to keep the original ByteDance-internal main branch intact for backward compatibility with existing users.

What's added

src/diffusers/
├── models/unets/unet_dreamlite.py            # DreamLiteUNetModel
├── pipelines/dreamlite/
│   ├── __init__.py
│   ├── pipeline_dreamlite.py                  # DreamLitePipeline (3-branch dual CFG)
│   ├── pipeline_dreamlite_mobile.py           # DreamLiteMobilePipeline (distilled)
│   └── pipeline_output.py
└── (registered in src/diffusers/__init__.py, models/__init__.py,
    pipelines/__init__.py, utils/dummy_*.py)

docs/source/en/api/pipelines/dreamlite.md
tests/pipelines/dreamlite/
├── test_pipeline_dreamlite.py
└── test_pipeline_dreamlite_mobile.py

Architecture highlights

DreamLiteUNetModel — UNet-based denoiser conditioned on Qwen3-VL text/vision embeddings.
DreamLitePipeline — runs 3 forward passes per step (text-cond / image-cond / uncond) and combines them with a dual-CFG schedule for high-fidelity text-to-image and image edit.
DreamLiteMobilePipeline — distilled single-pass variant; no CFG; designed for on-device inference. Pairs with AutoencoderTiny.
Both pipelines use FlowMatchEulerDiscreteScheduler.

Testing

Loading smoke test against carlofkl/DreamLite-base with revision="diffusers" — all 6 sub-modules resolve to the correct diffusers.* namespace.
Inference smoke test — generates a 1024×1024 image in ~0.6s/step on a single A800; output stats sane (std≈93, no NaN/Inf).
Standard pipeline tests in tests/pipelines/dreamlite/.

Before submitting

Did you read the contributor guideline?
Did you read our philosophy doc?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case. → Release DreamLite on Hugging Face ByteVisionLab/DreamLite#3
Did you make sure to update the documentation with your changes? Here are the documentation guidelines.
Did you write any new necessary tests?

Who can review?

cc @sayakpaul @yiyixuxu @DN6 — thanks in advance for the review!

Add ByteDance's DreamLite model family to diffusers. DreamLite is a UNet-based diffusion model that supports both text-to-image generation and reference-image editing through a shared 3-branch dual-CFG design. Two pipelines are shipped: * DreamLitePipeline - full 3-branch dual CFG (negative, reference, prompt); supports T2I and I2I editing at 1024x1024. * DreamLiteMobilePipeline - distilled single-branch variant for on-device inference; no CFG. New model code (all isolated under *_dreamlite.py / unet_dreamlite.py to avoid touching shared upstream files): * models/transformers/transformer_2d_dreamlite.py - DreamLite 2D transformer block. * models/unets/unet_dreamlite.py - DreamLiteUNetModel. * models/unets/unet_2d_blocks_dreamlite.py - DreamLite-specific down/up/mid blocks. * models/resnet_dreamlite.py - DreamLite ResNet variants. * models/attention_processor.py - add DreamLiteAttnProcessor2_0 (pure addition, no existing processor modified). Pipeline + tests + docs: * pipelines/dreamlite/{__init__.py, pipeline_dreamlite.py, pipeline_dreamlite_mobile.py, pipeline_output.py}. * tests/pipelines/dreamlite/{test_pipeline_dreamlite.py, test_pipeline_dreamlite_mobile.py} with the standard PipelineTesterMixin suite; setUp/tearDown auto-patches encode_prompt with a fake so MagicMock text encoders work without per-test boilerplate. * Skip 8 mixin tests that don't apply to DreamLite (MagicMock serialisation, custom attention processor, encode_prompt return shape, batch_size > 1 sweep), mirroring SD3 / Flux conventions. * docs/source/en/api/pipelines/dreamlite.md + _toctree.yml entry (alphabetically between DiT and EasyAnimate). * Register exports in 6 __init__.py files. Two real bugs surfaced by the mixin test suite are fixed in this commit: * num_images_per_prompt > 1: prompt_embeds and text_attention_mask are now repeated along the batch dimension in both pipelines' T2I and I2I branches before being passed to the UNet. * vae=None: __init__ now guards the encoder_block_out_channels lookup so encode_prompt can be tested in isolation per PipelineTesterMixin convention. SlowTests real-checkpoint resolution is set to 1024x1024 (the only size DreamLite is trained for). Test result: 27 passed, 50 skipped, 0 failed on CPU fast suite. make style && make quality: clean.

The `carlofkl/DreamLite-{base,mobile}` Hub repos host two flavours of the same checkpoint: * `main` branch - keeps `model_index.json` pointing at ByteDance's internal package path so the original (non-diffusers) reference code can still load these weights. * `diffusers` branch - rewrites the `unet` entry of `model_index.json` to `["diffusers", "DreamLiteUNetModel"]` so this integration loads correctly from `diffusers`. This commit pins every `from_pretrained(...)` call shipped with the diffusers integration (docs examples, pipeline docstrings, SlowTests) to `revision="diffusers"`. Local-override env vars (DREAMLITE_BASE_PATH / DREAMLITE_MOBILE_PATH) still bypass the revision pin.

…ts after rebase Mechanical changes after rebasing onto current `main`: * `pipeline_dreamlite.py::retrieve_timesteps` — re-synced from `diffusers.pipelines.flux.pipeline_flux.retrieve_timesteps` (PEP 604 type hints, expanded docstring, plus the new `accepts_timesteps` / `accept_sigmas` introspection guards). DreamLite's default code path uses `num_inference_steps` (uniform schedule) and never passes custom `timesteps` / `sigmas`, so the added guards are dead-code for this pipeline — behaviour is unchanged. * `dummy_pt_objects.py` / `dummy_torch_and_transformers_objects.py` — registered the dummy classes auto-generated by `make fix-copies` for `DreamLiteTransformer2DModel`, `DreamLiteUNetModel`, `DreamLitePipeline`, `DreamLiteMobilePipeline`, `DreamLitePipelineOutput`. Generated by `make fix-copies`. No hand edits.

HuggingFaceDocBuilderDev · 2026-05-27T23:22:11Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

…ing entries - Register DreamLiteAttnProcessor2_0 in docs/source/en/api/attnprocessor.md (fixes check_support_list.py). - Split combined 'height / width' and 'guidance_scale / image_guidance_scale' entries in the two pipeline docstrings; add a complete Args block to DreamLiteTransformer2DModel.forward (fixes check_forward_call_docstrings.py). No behavioral change.

Carlofkl · 2026-05-28T04:40:30Z

Hi @sayakpaul @yiyixuxu — pushed a small follow-up commit (032412c) that fixes the two check_repository_consistency failures from the previous run:

Registered DreamLiteAttnProcessor2_0 in docs/source/en/api/attnprocessor.md (was missing — check_support_list.py).
Split combined height / width and guidance_scale / image_guidance_scale docstring entries in both pipelines into separate lines, and added a complete Args block to DreamLiteTransformer2DModel.forward (was tripping check_forward_call_docstrings.py).

No behavioral change — docs/docstrings only. Verified both lints pass locally.

Whenever convenient, could you re-approve the workflows? Thanks!

Carlofkl · 2026-06-02T08:33:11Z

Hi @yiyixuxu @DN6 @sayakpaul — quick update: CI is now fully green
(all 15 checks passing) and I've just marked the PR as ready for review.
Whenever you have a moment to take a look, I'd really appreciate it. Thanks!

dg845 · 2026-06-11T04:44:36Z

@bot /style

github-actions · 2026-06-11T04:45:48Z

Style bot fixed some files and pushed the changes.

dg845 · 2026-06-11T06:51:08Z

+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+class DreamLiteUNetModel(UNet2DConditionModel):


suggestion: rather than inheriting directly from UNet2DConditionModel, DreamLiteUNetModel could inherit from the appropriate mixins (ModelMixin, ConfigMixin, AttentionMixin, etc.) and implement the forward logic directly. I think this would make the exact logic used by the DreamLite U-Net more clear and fits the single-file design better. CC @yiyixuxu

dg845 · 2026-06-11T07:02:45Z

+        original_forward = pipe.unet.forward
+        seen_batches = []
+
+        def spy_forward(*args, **kwargs):
+            x = args[0] if args else kwargs["sample"]
+            seen_batches.append(x.shape[0])
+            return original_forward(*args, **kwargs)
+
+        pipe.unet.forward = spy_forward
+        inputs = self.get_dummy_i2i_inputs(device)
+        inputs["num_inference_steps"] = 1
+        pipe(**inputs)


Suggested change

original_forward = pipe.unet.forward

seen_batches = []

def spy_forward(*args, **kwargs):

x = args[0] if args else kwargs["sample"]

seen_batches.append(x.shape[0])

return original_forward(*args, **kwargs)

pipe.unet.forward = spy_forward

inputs = self.get_dummy_i2i_inputs(device)

inputs["num_inference_steps"] = 1

pipe(**inputs)

inputs = self.get_dummy_i2i_inputs(device)

inputs["num_inference_steps"] = 1

with patch.object(pipe.unet, "forward", wraps=pipe.unet.forward) as spy:

pipe(**inputs)

seen_batches = [call.args[0].shape[0] for call in spy.call_args_list]

nit: I think using unittest.mock.patch here would work better than doing it manually (for example, it's more robust if an error occurs).

dg845 · 2026-06-11T07:05:12Z

Hi @claude, can you do a review for this PR?

github-actions · 2026-06-11T07:05:37Z

Claude finished @dg845's task in 4m 57s —— View job

Review of DreamLite Integration PR

Read repository review guidelines (AGENTS.md, models.md, pipelines.md, modular.md)
Read all source files under src/diffusers/
Read test files
Trace call paths and check for dead code
Post review feedback

Thanks for the work on this PR! I've done a thorough review of all the files under src/diffusers/ and the tests. There's a lot that looks good — the single-file model design, the dispatch_attention_fn usage, the # Copied from markers, and the test setup with a real tiny Qwen3-VL. Below are the issues I found, organized by severity.

Blocking Issues

1. _no_split_modules and _repeated_blocks are inherited from the parent but reference wrong class names
src/diffusers/models/unets/unet_dreamlite.py

DreamLiteUNetModel inherits UNet2DConditionModel's class attributes:

# From UNet2DConditionModel (parent):
_no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"]
_repeated_blocks = ["BasicTransformerBlock"]

But DreamLite uses BasicTransformerBlockDreamLite, ResnetBlock2DDreamLite, DreamLiteCrossAttnUpBlock2D, DreamLiteUpBlock2D. Since _no_split_modules is matched by __class__.__name__, none of the DreamLite modules will match, meaning device_map="auto" will have no no-split protection, and compile_repeated_blocks() will raise ValueError (no matching blocks). Override these on DreamLiteUNetModel:

_no_split_modules = [
    "BasicTransformerBlockDreamLite", "ResnetBlock2DDreamLite",
    "DreamLiteCrossAttnUpBlock2D", "DreamLiteUpBlock2D",
]
_repeated_blocks = ["BasicTransformerBlockDreamLite"]

2. Attention mask format will break non-SDPA backends (flash, sage, etc.)
src/diffusers/models/transformers/transformer_2d_dreamlite.py:549-555

The transformer converts the encoder attention mask to a dense additive float mask ((1 - mask) * -10000.0). Per the repo's attention guidelines in models.md, dense additive float masks hard-raise on flash / _flash_3 / _sage backends. Since DreamLite uses dispatch_attention_fn (which supports these backends), this limits backend compatibility.

Recommendation: Keep the mask as a bool tensor and let dispatch_attention_fn handle normalization via _normalize_attn_mask. If all sequences in the batch are the same length, set the mask to None (the Flux pattern). The DreamLiteAttnProcessor2_0 at unet_dreamlite.py:320-322 also runs prepare_attention_mask which produces an additive float mask — that should be updated to match.

3. Operator precedence bug in check_inputs
src/diffusers/pipelines/dreamlite/pipeline_dreamlite.py:372-377 (same in mobile)

if (
    height is not None
    and height % self.vae_scale_factor != 0
    or width is not None          # <-- `or` has lower precedence than `and`
    and width % self.vae_scale_factor != 0
):

Python's and binds tighter than or, so this parses as:

(height is not None and height % vae_scale_factor != 0) or (width is not None and width % vae_scale_factor != 0)

which is actually correct by accident — but it's a readability trap. Add explicit parentheses for each clause to make the intent unambiguous:

if (
    (height is not None and height % self.vae_scale_factor != 0)
    or (width is not None and width % self.vae_scale_factor != 0)
):

Non-Blocking Issues

4. Massive amount of dead code carried from parent / reference repo
Per the coding guidelines: "When porting from a research repo, delete training-time code paths, experimental flags, and ablation branches entirely — only keep the inference path you are actually integrating."

The following code paths are never exercised by DreamLite's default config or pipeline:

Dead code	Location	Reason
`DualTransformer2DModel` import + all `dual_cross_attention` branches	`unet_dreamlite.py:52`, every block class	`dual_cross_attention` defaults to `False`, never set `True`
`ada_norm_zero`, `ada_norm_single`, `ada_norm_continuous`, `layer_norm_i2vgen` branches	`transformer_2d_dreamlite.py` — ~80 lines across `__init__` and `forward`	DreamLite only uses `norm_type="layer_norm"`
`kernel="fir"` and `kernel="sde_vp"` paths in `ResnetBlock2DDreamLite`	`unet_dreamlite.py:201-213`	Default kernel is `None`
`use_sep_conv=False` branch in `ResnetBlock2DDreamLite`	`unet_dreamlite.py:167-168, 195`	Default is `True`
`GatedSelfAttentionDense` / `attention_type="gated"` branch	`transformer_2d_dreamlite.py:242-243`	`attention_type` defaults to `"default"`
`SinusoidalPositionalEmbedding` / `positional_embeddings` support	`transformer_2d_dreamlite.py:143-145`	Never set
`scale_shift_table` for PixArt-Alpha	`transformer_2d_dreamlite.py:247`	`norm_type` is never `ada_norm_single`

This is significant — BasicTransformerBlockDreamLite is essentially a full copy of BasicTransformerBlock with one added use_self_attention knob. This runs counter to the single-file model convention, which aims for minimal self-contained implementations. I'd strongly recommend stripping the unused branches (ada_norm variants, gated attention, dual cross attention, FIR/SDE_VP kernels) to keep only the code DreamLite actually exercises. This agrees with dg845's comment about inheriting from mixins directly rather than from UNet2DConditionModel.

5. DreamLiteTransformer2DModel is registered as a standalone model but isn't loadable independently
src/diffusers/models/transformers/transformer_2d_dreamlite.py

DreamLiteTransformer2DModel has ModelMixin + ConfigMixin + @register_to_config, which means it's a fully serializable model. But it's only ever used as an internal submodule inside the UNet blocks — it's never registered as a pipeline component or loaded via from_pretrained. The __init__.py registrations expose it at the top level (from diffusers import DreamLiteTransformer2DModel), which suggests it's a loadable model. This is misleading — either it should be a plain nn.Module (not a ModelMixin subclass) if it's only a building block, or the export should be removed from __init__.py.

6. vae_scale_factor computation differs from the pipeline convention
src/diffusers/pipelines/dreamlite/pipeline_dreamlite.py:204-207

if self.vae is not None and hasattr(self.vae.config, "encoder_block_out_channels"):
    self.vae_scale_factor = 2 ** (len(self.vae.config.encoder_block_out_channels) - 1)
else:
    self.vae_scale_factor = 8

Then on line 209:

self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)

The * 2 multiplication is unusual and undocumented. With the test's AutoencoderTiny(encoder_block_out_channels=(32, 32)), vae_scale_factor = 2^1 = 2, so VaeImageProcessor gets 4. With the real model (32, 32, 64), it would be 2^2 = 4, so VaeImageProcessor gets 8. Is the * 2 intentional to account for AutoencoderTiny's behavior? This should at minimum have a comment explaining why.

7. text_pad_embedding parameter — ephemeral or intentional?
Both pipelines accept text_pad_embedding in __call__ and encode_prompt, but the pipeline never generates or stores one internally. No test exercises this path. If this is an experimental/advanced feature, consider whether it should be deferred to a follow-up. If it's intentional, it needs a test.

8. Missing # Copied from on calculate_shift in pipeline_dreamlite.py
src/diffusers/pipelines/dreamlite/pipeline_dreamlite.py:82

The calculate_shift function is nearly identical to the one in pipeline_flux.py (differing only in max_shift default: 1.16 vs 1.15). The mobile pipeline correctly has # Copied from diffusers.pipelines.dreamlite.pipeline_dreamlite.calculate_shift, but the base pipeline's calculate_shift has no provenance marker. Since it's not identical to Flux's, it can't use # Copied from the Flux version — but it should either link to the source it was adapted from or note why 1.16 was chosen.

9. nn.Module.__init__(self) call bypasses MRO
src/diffusers/models/unets/unet_dreamlite.py:1639

nn.Module.__init__(self)

This deliberately skips UNet2DConditionModel.__init__ (with a comment explaining why), but it also skips ModelMixin.__init__, ConfigMixin.__init__, AttentionMixin.__init__, and all other mixin __init__ methods in the MRO. This could cause subtle issues with mixin state that expects initialization. This is another argument for dg845's suggestion to inherit from the mixins directly rather than from UNet2DConditionModel.

10. Test spy pattern could use unittest.mock.patch
tests/pipelines/dreamlite/test_pipeline_dreamlite.py:275-286

As dg845 noted, the manual forward-spy pattern is fragile. Using unittest.mock.patch.object(pipe.unet, "forward", wraps=pipe.unet.forward) is more robust. Same applies to the mobile test at test_pipeline_dreamlite_mobile.py:258-269.

Suggestions / Additional Info (Advisory — Dead Code Trace)

I traced the call path from both pipelines' __call__ through the UNet to identify likely dead code. Findings are advisory since the actual model config on the Hub may differ from defaults:

Parameters declared in UNet __init__ but never passed by the pipeline:

class_labels — never passed; all class_embed_type logic is dead under default config
cross_attention_kwargs — never passed; all cross_attention_kwargs forwarding in blocks is no-op
attention_mask — never passed to the UNet; only encoder_attention_mask is used
down_block_additional_residuals, mid_block_additional_residual — ControlNet residuals, not used

Methods on DreamLiteUNetModel inherited from parent that are never called:

enable_freeu / disable_freeu — FreeU is never configured by the pipeline
The is_freeu_enabled checks in up blocks (unet_dreamlite.py:1038-1043, 1198-1203, 1302-1307) are always false

text_pad_embedding in both pipelines — the parameter is accepted but never generated internally. The code path at pipeline_dreamlite.py:308-316 and pipeline_dreamlite_mobile.py:312-320 is dead unless an advanced user explicitly passes it.

Overall the PR is well-structured and the iteration with dg845 has resolved most earlier review feedback cleanly. The main remaining architectural decision is whether to strip the dead code inherited from the parent UNet (issues 4, 5, 9) — which aligns with dg845's suggestion to inherit from the mixins directly. The _no_split_modules issue (1) and the attention mask format (2) are the most impactful correctness issues.

dg845

Thanks for the changes! I think this PR is close to merge.

For the changes Claude suggested, I think they all look reasonable and in particular the blocking issues (1, 2, 3) and issue 9 (making sure that the mixin initializations aren't broken) look like the most important issues to address.

Also, can you look into the DreamLite pipeline test failures in this CI run? The issues look like they are coming from the Qwen3VLForConditionalGeneration text encoder, could this be due to a mismatch between the transformers version used in the CI and the version expected by the DreamLite code?

- Override _no_split_modules / _repeated_blocks on DreamLiteUNetModel with the actual DreamLite class names (BasicTransformerBlockDreamLite, ResnetBlock2DDreamLite, DreamLiteCrossAttnUpBlock2D, DreamLiteUpBlock2D) so device_map="auto" and compile_repeated_blocks() match correctly. - Keep attention masks as bool tensors in DreamLiteTransformer2DModel instead of converting them to dense additive float biases. The dense format hard-raises on flash / _flash_3 / _sage backends in dispatch_attention_fn (which requires dtype == torch.bool). - Add explicit parentheses around each clause in check_inputs's mixed and/or condition (both pipelines) for readability. - Replace nn.Module.__init__(self) with ModelMixin.__init__(self) in DreamLiteUNetModel.__init__ so mixin state (e.g. _gradient_checkpointing_func) is properly initialised. ConfigMixin / PushToHubMixin don't define their own __init__, so this covers the full chain without re-running UNet2DConditionModel.__init__.

Recent versions of Qwen3VLProcessor add an mm_token_type_ids output, and Qwen3VLModel.compute_3d_position_ids raises ValueError whenever multimodal inputs are present (image_grid_thw is not None) but mm_token_type_ids is None. encode_prompt previously forwarded only input_ids / attention_mask / pixel_values / image_grid_thw, dropping the new field and breaking the fast pipeline tests against transformers main. Switch to ``self.text_encoder(**tk_out, output_hidden_states=True)`` (matching NucleusMoEImagePipeline) so all processor outputs are forwarded automatically and future additions don't regress this path.

Carlofkl · 2026-06-11T15:16:24Z

Thanks @dg845! Pushed two commits:

6653dd491 — addresses issues 1, 2, 3, and 9.
a02bb23be — fixes the CI failures.

The CI failures were caused by recent versions of Qwen3VLProcessor emitting a new mm_token_type_ids field, which Qwen3VLModel.compute_3d_position_ids now requires whenever image_grid_thw is provided. Switched both pipelines to self.text_encoder(**tk_out, output_hidden_states=True) (matching NucleusMoEImagePipeline) so all processor outputs are forwarded automatically.

Could you re-approve the workflow run when you get a chance?

dg845 · 2026-06-12T01:09:26Z

@bot /style

github-actions · 2026-06-12T01:10:53Z

Style bot fixed some files and pushed the changes.

dg845

Thanks for your work on this PR! Left some small comments on the docs/examples.

- Replace broken cat.png URL in editing examples (both base and mobile) with the standard `huggingface/documentation-images` source used elsewhere in the diffusers docs. - Promote the recommended guidance_scale=3.5 / image_guidance_scale=1.5 to the default values of DreamLitePipeline.__call__, and drop the now-redundant explicit args from the docs examples. - Switch the EXAMPLE_DOC_STRING examples in both pipelines from torch.float16 to torch.bfloat16 for consistency with the rest of the docs.

Carlofkl · 2026-06-12T05:38:00Z

Thanks @dg845! Pushed f38937d97 addressing all five nits:

1 & 3. Replaced the broken cat.png URL with the standard huggingface/documentation-images/.../diffusers/cat.png (verified 200 OK).

Promoted guidance_scale=3.5 / image_guidance_scale=1.5 to defaults on DreamLitePipeline.__call__ and removed the now-redundant explicit args from the docs examples.

4 & 5. Switched the EXAMPLE_DOC_STRING examples to torch.bfloat16 in both pipelines.

dg845 · 2026-06-12T06:04:59Z

Merging as the CI failures are unrelated.

) * feat(pipelines): add DreamLite text-to-image and image-edit pipelines Add ByteDance's DreamLite model family to diffusers. DreamLite is a UNet-based diffusion model that supports both text-to-image generation and reference-image editing through a shared 3-branch dual-CFG design. Two pipelines are shipped: * DreamLitePipeline - full 3-branch dual CFG (negative, reference, prompt); supports T2I and I2I editing at 1024x1024. * DreamLiteMobilePipeline - distilled single-branch variant for on-device inference; no CFG. New model code (all isolated under *_dreamlite.py / unet_dreamlite.py to avoid touching shared upstream files): * models/transformers/transformer_2d_dreamlite.py - DreamLite 2D transformer block. * models/unets/unet_dreamlite.py - DreamLiteUNetModel. * models/unets/unet_2d_blocks_dreamlite.py - DreamLite-specific down/up/mid blocks. * models/resnet_dreamlite.py - DreamLite ResNet variants. * models/attention_processor.py - add DreamLiteAttnProcessor2_0 (pure addition, no existing processor modified). Pipeline + tests + docs: * pipelines/dreamlite/{__init__.py, pipeline_dreamlite.py, pipeline_dreamlite_mobile.py, pipeline_output.py}. * tests/pipelines/dreamlite/{test_pipeline_dreamlite.py, test_pipeline_dreamlite_mobile.py} with the standard PipelineTesterMixin suite; setUp/tearDown auto-patches encode_prompt with a fake so MagicMock text encoders work without per-test boilerplate. * Skip 8 mixin tests that don't apply to DreamLite (MagicMock serialisation, custom attention processor, encode_prompt return shape, batch_size > 1 sweep), mirroring SD3 / Flux conventions. * docs/source/en/api/pipelines/dreamlite.md + _toctree.yml entry (alphabetically between DiT and EasyAnimate). * Register exports in 6 __init__.py files. Two real bugs surfaced by the mixin test suite are fixed in this commit: * num_images_per_prompt > 1: prompt_embeds and text_attention_mask are now repeated along the batch dimension in both pipelines' T2I and I2I branches before being passed to the UNet. * vae=None: __init__ now guards the encoder_block_out_channels lookup so encode_prompt can be tested in isolation per PipelineTesterMixin convention. SlowTests real-checkpoint resolution is set to 1024x1024 (the only size DreamLite is trained for). Test result: 27 passed, 50 skipped, 0 failed on CPU fast suite. make style && make quality: clean. * docs+tests(pipelines/dreamlite): pin Hub repos to `diffusers` branch The `carlofkl/DreamLite-{base,mobile}` Hub repos host two flavours of the same checkpoint: * `main` branch - keeps `model_index.json` pointing at ByteDance's internal package path so the original (non-diffusers) reference code can still load these weights. * `diffusers` branch - rewrites the `unet` entry of `model_index.json` to `["diffusers", "DreamLiteUNetModel"]` so this integration loads correctly from `diffusers`. This commit pins every `from_pretrained(...)` call shipped with the diffusers integration (docs examples, pipeline docstrings, SlowTests) to `revision="diffusers"`. Local-override env vars (DREAMLITE_BASE_PATH / DREAMLITE_MOBILE_PATH) still bypass the revision pin. * chore(pipelines/dreamlite): sync `# Copied from` blocks + dummy objects after rebase Mechanical changes after rebasing onto current `main`: * `pipeline_dreamlite.py::retrieve_timesteps` — re-synced from `diffusers.pipelines.flux.pipeline_flux.retrieve_timesteps` (PEP 604 type hints, expanded docstring, plus the new `accepts_timesteps` / `accept_sigmas` introspection guards). DreamLite's default code path uses `num_inference_steps` (uniform schedule) and never passes custom `timesteps` / `sigmas`, so the added guards are dead-code for this pipeline — behaviour is unchanged. * `dummy_pt_objects.py` / `dummy_torch_and_transformers_objects.py` — registered the dummy classes auto-generated by `make fix-copies` for `DreamLiteTransformer2DModel`, `DreamLiteUNetModel`, `DreamLitePipeline`, `DreamLiteMobilePipeline`, `DreamLitePipelineOutput`. Generated by `make fix-copies`. No hand edits. * docs(dreamlite): register attention processor + split combined docstring entries - Register DreamLiteAttnProcessor2_0 in docs/source/en/api/attnprocessor.md (fixes check_support_list.py). - Split combined 'height / width' and 'guidance_scale / image_guidance_scale' entries in the two pipeline docstrings; add a complete Args block to DreamLiteTransformer2DModel.forward (fixes check_forward_call_docstrings.py). No behavioral change. * refactor(dreamlite): address review feedback from #13815 - Inline the down/up block factories and define DreamLiteCrossAttn{,NoSelfAttn}{Down,Up}Block2D directly (review #1, #2) - Rename DownBlock2DDreamLite/UpBlock2DDreamLite to DreamLiteDownBlock2D/DreamLiteUpBlock2D to match diffusers naming conventions (review #3, #4) - Merge unet_2d_blocks_dreamlite.py into unet_dreamlite.py to mirror recent transformer model files (review #5) - Wire max_sequence_length into the tokenizer call for generate mode (review #6) - Replace hard-coded drop_idx values (64/34) with self.prompt_template_encode_*_start_idx attributes plus a comment explaining how the offsets are derived (review #7, #8) - Drop the manual Image.resize call and rely on VaeImageProcessor's LANCZOS default in preprocess(image, height, width) (review #9) - Use self.guidance_scale / self.image_guidance_scale properties in the CFG combine instead of the underscore-prefixed attributes (review #10, #11) - Inline retrieve_latents / retrieve_timesteps / calculate_shift in the mobile pipeline with `# Copied from` markers, removing the cross-pipeline imports (review #12) - Add `# Copied from` marker to _extract_masked_hidden in the mobile pipeline (review #13) * refactor(dreamlite): address dg845 follow-up review - Merge resnet_dreamlite.py (DepthwiseSeparableConv + ResnetBlock2DDreamLite) into unet_dreamlite.py and delete the standalone module (review #1) - Move DreamLiteAttnProcessor2_0 from attention_processor.py into unet_dreamlite.py to keep all DreamLite-specific code in one place; update docs autodoc reference accordingly (review #2) - Drop the PyTorch 2.0 hasattr/ImportError check in DreamLiteAttnProcessor2_0.__init__ (diffusers already requires torch>=2.0; matches Wan deprecation) (review #3) - Drop the deprecated `scale` argument handling from DreamLiteAttnProcessor2_0.__call__ (new model, no legacy callers) (review #4) - Switch SDPA call to dispatch_attention_fn so all diffusers attention backends (FlashAttention, FlashAttention-3, sageattention, etc.) are selectable (review #5) - Rename block dispatch keys in _get_{down,mid,up}_block_dreamlite to match the Python class names (DreamLiteCrossAttn{Down,Up}Block2D / DreamLiteCrossAttnNoSelfAttn{Down,Up}Block2D / DreamLiteUNetMidBlock2DCrossAttn / DreamLite{Down,Up}Block2D); default down/up/mid block_types in DreamLiteUNetModel and the test fixtures are updated to the new keys (review #6, #7); the carlofkl/DreamLite-{base,mobile} (diffusers branch) Hub configs are being updated in lock-step - Localize retrieve_latents inside pipeline_dreamlite.py with a `# Copied from` marker, removing the cross-pipeline import; mirrors the mobile pipeline (review #8) - Add a check_inputs() method to both DreamLitePipeline and DreamLiteMobilePipeline (mobile uses `# Copied from`); call it from __call__; pulls the image-type validation out of prepare_image_latents and adds prompt-type and h/w-divisibility checks (review #9) * fix(dreamlite): correct Q/K/V layout for dispatch_attention_fn dispatch_attention_fn expects (batch, seq, heads, head_dim) and handles the transpose internally; the previous code passed (batch, heads, seq, head_dim), which collided with the dispatch's internal transpose and broke inference (RuntimeError: tensor size mismatch at non-singleton dimension 1). * test(dreamlite): swap MagicMock for tiny real Qwen3-VL fixture Address dg845's review: rebuild the DreamLite fast-test fixture around a real (tiny) Qwen3VLForConditionalGeneration + Qwen3VLProcessor so the standard PipelineTesterMixin save/load, dtype, and offload tests run end-to-end against the actual encode_prompt code path. Override DreamLiteUNetModel.set_default_attn_processor to reinstall the GQA processor so mixin utilities that round-trip through it keep working. * Apply style fixes * fix(dreamlite): address blocking review issues from #13815 - Override _no_split_modules / _repeated_blocks on DreamLiteUNetModel with the actual DreamLite class names (BasicTransformerBlockDreamLite, ResnetBlock2DDreamLite, DreamLiteCrossAttnUpBlock2D, DreamLiteUpBlock2D) so device_map="auto" and compile_repeated_blocks() match correctly. - Keep attention masks as bool tensors in DreamLiteTransformer2DModel instead of converting them to dense additive float biases. The dense format hard-raises on flash / _flash_3 / _sage backends in dispatch_attention_fn (which requires dtype == torch.bool). - Add explicit parentheses around each clause in check_inputs's mixed and/or condition (both pipelines) for readability. - Replace nn.Module.__init__(self) with ModelMixin.__init__(self) in DreamLiteUNetModel.__init__ so mixin state (e.g. _gradient_checkpointing_func) is properly initialised. ConfigMixin / PushToHubMixin don't define their own __init__, so this covers the full chain without re-running UNet2DConditionModel.__init__. * fix(dreamlite): forward all processor outputs to Qwen3VL text encoder Recent versions of Qwen3VLProcessor add an mm_token_type_ids output, and Qwen3VLModel.compute_3d_position_ids raises ValueError whenever multimodal inputs are present (image_grid_thw is not None) but mm_token_type_ids is None. encode_prompt previously forwarded only input_ids / attention_mask / pixel_values / image_grid_thw, dropping the new field and breaking the fast pipeline tests against transformers main. Switch to ``self.text_encoder(**tk_out, output_hidden_states=True)`` (matching NucleusMoEImagePipeline) so all processor outputs are forwarded automatically and future additions don't regress this path. * Apply style fixes * docs(dreamlite): address final review nits from #13815 - Replace broken cat.png URL in editing examples (both base and mobile) with the standard `huggingface/documentation-images` source used elsewhere in the diffusers docs. - Promote the recommended guidance_scale=3.5 / image_guidance_scale=1.5 to the default values of DreamLitePipeline.__call__, and drop the now-redundant explicit args from the docs examples. - Switch the EXAMPLE_DOC_STRING examples in both pipelines from torch.float16 to torch.bfloat16 for consistency with the rest of the docs. --------- Co-authored-by: YiYi Xu <[email protected]> Co-authored-by: Sayak Paul <[email protected]> Co-authored-by: dg845 <[email protected]> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Carlofkl added 3 commits May 27, 2026 11:38

github-actions Bot added size/L PR with diff > 200 LOC documentation Improvements or additions to documentation models tests utils pipelines and removed size/L PR with diff > 200 LOC labels May 27, 2026

Carlofkl mentioned this pull request May 27, 2026

Release DreamLite on Hugging Face ByteVisionLab/DreamLite#3

Open

Merge branch 'main' into feature/dreamlite-integration

0b7d747

github-actions Bot added the size/L PR with diff > 200 LOC label May 27, 2026

Carlofkl marked this pull request as ready for review May 31, 2026 14:18

Merge branch 'main' into feature/dreamlite-integration

7d9bd46

sayakpaul requested review from dg845 and yiyixuxu June 2, 2026 08:34