diff --git a/.cursor/commands/scaffold.md b/.cursor/commands/scaffold.md new file mode 100644 index 000000000000..20ae8700cbe2 --- /dev/null +++ b/.cursor/commands/scaffold.md @@ -0,0 +1,56 @@ +# /scaffold — first contribution on the REAL library checkout + +Usage: `/scaffold ` + +This workspace is `huggingface/diffusers` (fork). Overlay kit is `ramp-kit/` +(cloned from alex-16moro/diffuser_agent). Do not overwrite `.ai/` or root `AGENTS.md`. + +`$1` = component (`scheduler`). `$2` = PascalCase name without suffix (`EulerLite` +→ `EulerLiteScheduler`, `scheduling_euler_lite.py`). + +## 0. Ground first + +Prefer **this checkout's source** over memory and over MCP: + +- `src/diffusers/schedulers/scheduling_euler_discrete.py` +- `src/diffusers/schedulers/scheduling_ddpm.py` + +Those files are the contract. `ramp-kit/conventions/rules.yaml` / the gate is +authoritative if anything disagrees (the philosophy doc is stale). + +Optional docs CLI (MCP is opt-in; `.cursor/mcp.json` is empty by default): + +```bash +python3 ramp-kit/tools/docs_mcp_server.py --query "scheduler set_timesteps step SchedulerMixin register_to_config" +``` + +Cite provenance if you run that. Do not read or copy +`ramp-kit/examples/candidate_scheduler/` into the new files. + +## 1. Paths (library layout, not the kit stand-in) + +- Implementation: `src/diffusers/schedulers/scheduling_euler_lite.py` +- Test: `tests/schedulers/test_scheduling_euler_lite.py` + +## 2. Copy from the overlay templates + +- `ramp-kit/templates/scheduler/scheduling_TEMPLATE.py` → implementation +- Rename `TemplateScheduler` → `$2Scheduler` +- Leave `TODO(engineer)` in `step`. Do not invent Euler math. + +## 3. TEST001 + +Copy `ramp-kit/tests/_templates/scheduler_test.py`. Set `TARGET` and `CLASS`. +The test must mention `set_timesteps` and `step`, and (TEST002) include +assertions, same-seed determinism, and shape/dtype checks. + +## 4. Gate (file-scoped — do not `--all` this library) + +```bash +python3 ramp-kit/tools/convention_check.py ramp-kit/examples/candidate_scheduler +python3 ramp-kit/tools/convention_check.py src/diffusers/schedulers/scheduling_euler_lite.py +python3 -m unittest tests.schedulers.test_scheduling_euler_lite -v +``` + +Catch-early uses the kit fixture path. Stop at 0 blocking. Do not fabricate numerics. +Do not open a PR against huggingface/diffusers — PR this fork. diff --git a/.cursor/commands/search-docs.md b/.cursor/commands/search-docs.md new file mode 100644 index 000000000000..9bfe92bf0ebb --- /dev/null +++ b/.cursor/commands/search-docs.md @@ -0,0 +1,24 @@ +# /search-docs — library docs in THIS checkout (optional) + +Usage: `/search-docs ` + +Prefer source + gate first (MCP is opt-in; `.cursor/mcp.json` is empty): + +- `src/diffusers/schedulers/scheduling_euler_discrete.py` +- `src/diffusers/schedulers/scheduling_ddpm.py` +- `ramp-kit/conventions/rules.yaml` + +Then, optional CLI (same server as MCP, no OAuth): + +```bash +python3 ramp-kit/tools/docs_mcp_server.py --query $1 +``` + +If `$1` is empty: + +```bash +python3 ramp-kit/tools/docs_mcp_server.py --query "scheduler set_timesteps step SchedulerMixin register_to_config" +``` + +Provenance should say `diffusers checkout`, not `bundled snapshot`. If the +MCP tool `search_docs` is in your list, call that instead. diff --git a/.cursor/environment.json b/.cursor/environment.json new file mode 100644 index 000000000000..1f7395a6a132 --- /dev/null +++ b/.cursor/environment.json @@ -0,0 +1,6 @@ +{ + "install": "if [ ! -d ramp-kit/.git ]; then git clone --depth 1 https://github.com/alex-16moro/diffuser_agent.git ramp-kit; fi && pip install -r ramp-kit/requirements.txt && python3 ramp-kit/tools/docs_mcp_server.py --selftest", + "repositoryDependencies": [ + "github.com/alex-16moro/diffuser_agent" + ] +} diff --git a/.cursor/hooks.json b/.cursor/hooks.json new file mode 100644 index 000000000000..033390d6887e --- /dev/null +++ b/.cursor/hooks.json @@ -0,0 +1,11 @@ +{ + "version": 1, + "hooks": { + "afterFileEdit": [ + { + "command": ".cursor/hooks/convention-gate.sh", + "timeout": 20 + } + ] + } +} diff --git a/.cursor/hooks/convention-gate.sh b/.cursor/hooks/convention-gate.sh new file mode 100755 index 000000000000..63a8672bdf34 --- /dev/null +++ b/.cursor/hooks/convention-gate.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# Cursor afterFileEdit hook. Project hooks run from the repo root. +# Reads the afterFileEdit JSON payload on stdin and runs the same +# convention gate as CI (`tools/convention_check.py`). +set -euo pipefail +ROOT="$(cd "$(dirname "$0")/../.." && pwd)" +exec python3 "$ROOT/.cursor/hooks/convention_gate.py" diff --git a/.cursor/hooks/convention_gate.py b/.cursor/hooks/convention_gate.py new file mode 100644 index 000000000000..d273d6eb612b --- /dev/null +++ b/.cursor/hooks/convention_gate.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +"""afterFileEdit hook — run the same convention gate CI uses. + +Cursor sends JSON on stdin: + {"file_path": "", "edits": [...]} + +This is a notification (the write already happened). Findings are printed +to stderr so the agent can fix them immediately. We always exit 0 so a +red gate never crashes the editor loop; blocking findings still show up. +""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +_kit_tools = None +for _cand in (ROOT / "ramp-kit" / "tools", ROOT / "tools"): + if (_cand / "convention_check.py").is_file(): + _kit_tools = _cand + break +sys.path.insert(0, str(_kit_tools or (ROOT / "tools"))) + +from convention_check import check_file, load_rules, render_human # noqa: E402 + +# Don't fire the gate on the kit's own machinery or the teaching fixture +# (the fixture is scanned explicitly by `make demo`). +_SKIP_SUBSTR = ( + "/tools/", + "/.cursor/", + "/_templates/", + "/templates/", + "/__pycache__/", + "/ramp-kit/examples/candidate_scheduler/", +) + + +def main() -> int: + raw = sys.stdin.read() + if not raw.strip(): + return 0 + try: + payload = json.loads(raw) + except json.JSONDecodeError: + return 0 + file_path = payload.get("file_path") or "" + path = Path(file_path) + if path.suffix != ".py" or not path.is_file(): + return 0 + posix = path.as_posix() + if any(s in posix for s in _SKIP_SUBSTR): + return 0 + findings = check_file(path, load_rules()) + sys.stderr.write(render_human(findings, 1)) + sys.stderr.flush() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.cursor/mcp-diffusers-docs.py b/.cursor/mcp-diffusers-docs.py new file mode 100755 index 000000000000..2b6aa4f7f379 --- /dev/null +++ b/.cursor/mcp-diffusers-docs.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Launch diffusers-docs MCP from any cwd. + +Cloud Agent stdio cannot set `cwd` and does not expand `${workspaceFolder}`. +This launcher walks cwd, git root, /workspace, and its own path until it finds +`tools/docs_mcp_server.py`, then execs it. Desktop and Cloud then share one +command: `python3 -u .cursor/mcp-diffusers-docs.py`. +""" +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + + +def _candidates(): + out: list[Path] = [] + here = Path(__file__).resolve().parent + out.append(here.parent) # repo root when this file lives in .cursor/ + out.append(here.parent / "ramp-kit") + cwd = Path.cwd().resolve() + out.extend([cwd, cwd / "ramp-kit", *cwd.parents]) + for key in ("CURSOR_PROJECT_DIR", "CURSOR_WORKSPACE", "WORKSPACE"): + val = os.environ.get(key) + if val: + out.append(Path(val)) + out.append(Path("/workspace")) + try: + top = subprocess.check_output( + ["git", "rev-parse", "--show-toplevel"], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + if top: + out.append(Path(top)) + except (OSError, subprocess.CalledProcessError): + pass + seen: set[Path] = set() + for raw in out: + try: + path = raw.resolve() + except OSError: + continue + if path in seen: + continue + seen.add(path) + yield path + + +def main(argv: list[str] | None = None) -> int: + extra = [a for a in (argv if argv is not None else sys.argv[1:]) if a != "--serve"] + for root in _candidates(): + script = root / "tools" / "docs_mcp_server.py" + if script.is_file(): + os.chdir(root) + sys.argv = [str(script), "--serve", *extra] + # runpy would work; exec keeps stdin/stdout as the MCP pipes. + os.execv(sys.executable, [sys.executable, "-u", str(script), "--serve", *extra]) + sys.stderr.write( + "diffusers-docs MCP: could not find tools/docs_mcp_server.py " + f"(cwd={Path.cwd()}). From the repo root run: " + "python3 -u .cursor/mcp-diffusers-docs.py\n" + ) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.cursor/mcp-diffusers-docs.sh b/.cursor/mcp-diffusers-docs.sh new file mode 100755 index 000000000000..852a35067450 --- /dev/null +++ b/.cursor/mcp-diffusers-docs.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +# Thin wrapper around the Python launcher (Cloud dashboard can use either). +set -euo pipefail +DIR="$(cd "$(dirname "$0")" && pwd)" +exec python3 -u "$DIR/mcp-diffusers-docs.py" --serve "$@" diff --git a/.cursor/mcp.json b/.cursor/mcp.json new file mode 100644 index 000000000000..da39e4ffafe8 --- /dev/null +++ b/.cursor/mcp.json @@ -0,0 +1,3 @@ +{ + "mcpServers": {} +} diff --git a/.cursor/mcp.optional.json b/.cursor/mcp.optional.json new file mode 100644 index 000000000000..ba79ce1a43be --- /dev/null +++ b/.cursor/mcp.optional.json @@ -0,0 +1,15 @@ +{ + "mcpServers": { + "diffusers-docs": { + "type": "stdio", + "command": "python3", + "args": ["-u", "/workspace/.cursor/mcp-diffusers-docs.py"], + "env": { + "PYTHONUNBUFFERED": "1" + } + }, + "huggingface": { + "url": "https://huggingface.co/mcp" + } + } +} diff --git a/.cursor/rules/00-conventions.mdc b/.cursor/rules/00-conventions.mdc new file mode 100644 index 000000000000..44cef080a400 --- /dev/null +++ b/.cursor/rules/00-conventions.mdc @@ -0,0 +1,38 @@ +--- +description: diffusers contribution conventions enforced by the Ramp Kit. Applied to all edits. +alwaysApply: true +--- + + +# diffusers conventions (registry v0.7.0) + +You are contributing to `huggingface/diffusers`. Follow these conventions. +They mirror the project's own `.ai/` rules and are enforced by +`tools/convention_check.py` — the same gate that runs in CI, so producing +code that violates a **block** rule will fail the build. + +Use the `diffusers-docs` MCP tool (`search_docs`) to ground answers in the +library's current docs before scaffolding. If that tool is **not** in your +tool list (Cloud Agents skip project `.cursor/mcp.json` unless the launch +MCP dropdown has stdio `python3 -u .cursor/mcp-diffusers-docs.py`), run +`/search-docs `, the `search-docs` skill, or +`python3 tools/docs_mcp_server.py --query "..."` — same server. +Do not pull in code or context from outside this repo's approved +boundaries (see `.cursorignore`). + +## Blocking conventions (must satisfy) +- **REPRO001 — Randomness threads through a generator, never global RNG.** Use `randn_tensor(shape, generator=generator, device=..., dtype=...)` and thread `generator` through the call chain. +- **DEVICE001 — No hardcoded CUDA/device placement.** Respect the module's existing device: `.to(sample.device)` or accept a `device` argument. +- **DEPR001 — No deprecated/moved import paths.** Use the current import path shown in the check output. +- **MUT001 — No mutable default arguments.** Default to None and initialise inside the body. + +## Advisory conventions (should satisfy) +- **DEPR002 — Avoid deprecated kwargs (verify per release).** Prefer the newer kwarg, but confirm against your pinned diffusers version. +- **COPY001 — # Copied from markers are well-formed.** Format: `# Copied from diffusers.. with A->B`. +- **LOG001 — Library code logs, it does not print().** Use `logger = logging.get_logger(__name__)` and `logger.info(...)`. +- **DOC001 — Public methods have docstrings.** Add a Google-style docstring with Args/Returns. +- **CUST001 — No debugging leftovers committed.** Remove breakpoint()/pdb before committing. + +Component-specific rules auto-attach when you open a matching file +(e.g. a scheduler). Prefer copying an existing in-repo example with a +`# Copied from` marker over inventing a new pattern. diff --git a/.cursor/rules/10-scheduler.mdc b/.cursor/rules/10-scheduler.mdc new file mode 100644 index 000000000000..96a442a25690 --- /dev/null +++ b/.cursor/rules/10-scheduler.mdc @@ -0,0 +1,35 @@ +--- +description: Extra diffusers conventions for schedulers. +globs: ["src/diffusers/schedulers/scheduling_*.py", "examples/**/scheduling_*.py", "tests/schedulers/test_scheduling_*.py"] +alwaysApply: false +--- + + +# Scheduler conventions (auto-attached to scheduler files) + +### SCHED001 — Schedulers inherit SchedulerMixin and ConfigMixin (block, owner: architect) +Schedulers are swappable via ConfigMixin.from_config and serializable via SchedulerMixin. A scheduler that skips these can't be loaded, saved, or swapped like every other scheduler — it breaks the pipeline contract. +*How:* Declare `class XScheduler(SchedulerMixin, ConfigMixin):`. + +### SCHED002 — Schedulers implement the step() / set_timesteps() contract (block, owner: architect) +Every denoising loop calls set_timesteps(...) once, then step(...) each iteration. VERIFIED against current source: DDPMScheduler and EulerDiscreteScheduler both define `set_timesteps(self, num_inference_steps, device=None, ...)` and `step(self, model_output, timestep, sample, generator=None, return_dict=True)`. The philosophy doc's older "set_num_inference_steps" name is stale; the code is canonical. +*How:* Implement `set_timesteps(self, num_inference_steps, device=None)` and `step(self, model_output, timestep, sample, generator=None, return_dict=True) -> SchedulerOutput`. + +### SCHED003 — Scheduler __init__ is decorated with @register_to_config (block, owner: architect) +@register_to_config captures constructor args into the config so from_pretrained / save_pretrained round-trip correctly. Without it, config is empty and the scheduler cannot be reconstructed from the Hub. +*How:* Add `@register_to_config` directly above `def __init__`. + +### IMPORT001 — Schedulers stay self-contained (no heavy util imports) (warn, owner: architect) +The single-file policy keeps schedulers readable and swappable. Pulling in pipeline internals or deep util chains couples them to the rest of the library and violates the philosophy reviewers enforce. +*How:* Copy the small helper in-file with a `# Copied from` marker instead of importing deep internals. + +### TEST001 — New schedulers/models ship with a matching test file (block, owner: qa) +"No quality testing = no merge." A new scheduler with no test cannot move toward deployment. Presence is not enough: the test must actually exercise set_timesteps and step, otherwise a dummy file would satisfy the gate. +*How:* Generate a contract test from tests/_templates/scheduler_test.py that asserts set_timesteps and step. + +### TEST002 — Scheduler tests assert behaviour, not just presence (block, owner: qa) +Presence of a test file is not enough. A test_* function with zero assertions is a dummy. New scheduler tests must include a same-seed determinism assertion and a shape/dtype assertion so CI actually pins the contract TEST001 only names. +*How:* Each test_* function must assert. Include same-seed determinism (torch.equal) and shape plus dtype checks (see tests/_templates/scheduler_test.py). + +Reference implementation: `examples/scaffolded_scheduler/scheduling_ddpm_lite.py`. +Start from the scaffold: `/scaffold scheduler ` — do not start from a blank file. diff --git a/.cursor/skills/search-docs/SKILL.md b/.cursor/skills/search-docs/SKILL.md new file mode 100644 index 000000000000..75fd9c762de6 --- /dev/null +++ b/.cursor/skills/search-docs/SKILL.md @@ -0,0 +1,15 @@ +--- +name: search-docs +description: Optional docs CLI for this diffusers checkout. Prefer scheduler source + the gate; MCP is opt-in. +--- + +# Search this library's docs + +```bash +python3 ramp-kit/tools/docs_mcp_server.py --query "" +``` + +Cite provenance (`diffusers checkout` vs bundled snapshot). Prefer reading +`src/diffusers/schedulers/scheduling_euler_discrete.py` and +`scheduling_ddpm.py` first. `ramp-kit/conventions/rules.yaml` is the gate. +`.cursor/mcp.json` is empty by default. diff --git a/.cursorignore b/.cursorignore new file mode 100644 index 000000000000..5e38f392af21 --- /dev/null +++ b/.cursorignore @@ -0,0 +1,10 @@ +# Approved context boundary for the overlay agent. + +# Known-bad teaching fixture (in the cloned kit). Gate it; do not copy it. +ramp-kit/examples/candidate_scheduler/ + +.git/ +**/__pycache__/ +**/*.pyc +.venv/ +ramp-kit/.git/ diff --git a/.gitignore b/.gitignore index 7b156e460abf..befbeba85af9 100644 --- a/.gitignore +++ b/.gitignore @@ -125,8 +125,7 @@ dmypy.json .vs .vscode -# Cursor -.cursor +# Cursor — upstream ignored this; overlay tracks .cursor (OVERLAY.md) # Pycharm .idea @@ -185,4 +184,6 @@ wandb # AI agent generated symlinks /.agents/skills -/.claude/skills \ No newline at end of file +/.claude/skills +# Ramp Kit overlay clone (see OVERLAY.md) +ramp-kit/ diff --git a/OVERLAY.md b/OVERLAY.md new file mode 100644 index 000000000000..ec43dcbd8049 --- /dev/null +++ b/OVERLAY.md @@ -0,0 +1,36 @@ +# Ramp Kit overlay (customer convention-as-code) + +This checkout is a **fork of huggingface/diffusers**. The overlay lives in a +separate repo: [alex-16moro/diffuser_agent](https://github.com/alex-16moro/diffuser_agent). + +Cloud Agent install clones it to `ramp-kit/` (gitignored). Do not edit +upstream `AGENTS.md` / `.ai/` — those stay Hugging Face's. + +PRs from this overlay are **fork-demo only, not for upstream**. Keep them draft +and titled `[fork demo — not for upstream]`. Overlay clearance is the customer +gate (`ramp-kit/tools/convention_check.py` on the new file). Upstream GitHub +Actions may go red; that is expected — we do not claim Hugging Face's CI. + +## Grounding (default path — no MCP) + +1. Read `src/diffusers/schedulers/scheduling_euler_discrete.py` and + `scheduling_ddpm.py` (code beats the philosophy doc). +2. Run the gate. `ramp-kit/conventions/rules.yaml` is authoritative. +3. Optional CLI docs query (same server as MCP, no OAuth): + +```bash +python3 ramp-kit/tools/docs_mcp_server.py --query "scheduler set_timesteps step SchedulerMixin register_to_config" +``` + +`.cursor/mcp.json` is **empty by default** so Cloud launches do not hit Hub +OAuth or stdio cwd failures. Opt-in servers: `.cursor/mcp.optional.json`. + +## Demo (Cloud Agent) + +1. On the **kit** first: `make demo-maintain` — one YAML edit, five surfaces. +2. Then launch on **this fork** (`alex-16moro/diffusers`), overlay branch. +3. Scaffold writes `src/diffusers/schedulers/scheduling_euler_lite.py` here. + Gate: `python3 ramp-kit/tools/convention_check.py ` (never `--all`). +4. Open the PR **on this fork**, not on huggingface/diffusers. + +Catch-early fixture: `ramp-kit/examples/candidate_scheduler` — do not copy it. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2825e9888c98..b7f45c83aa0d 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -447,6 +447,7 @@ "EntropyBoundSchedulerOutput", "EulerAncestralDiscreteScheduler", "EulerDiscreteScheduler", + "EulerLiteScheduler", "FlowMapEulerDiscreteScheduler", "FlowMatchEulerDiscreteScheduler", "FlowMatchHeunDiscreteScheduler", @@ -1325,6 +1326,7 @@ EntropyBoundSchedulerOutput, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, + EulerLiteScheduler, FlowMapEulerDiscreteScheduler, FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler, diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index c0e46ef445df..b803cbcb8297 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -60,6 +60,7 @@ _import_structure["scheduling_entropy_bound"] = ["EntropyBoundScheduler", "EntropyBoundSchedulerOutput"] _import_structure["scheduling_euler_ancestral_discrete"] = ["EulerAncestralDiscreteScheduler"] _import_structure["scheduling_euler_discrete"] = ["EulerDiscreteScheduler"] + _import_structure["scheduling_euler_lite"] = ["EulerLiteScheduler"] _import_structure["scheduling_flow_map_euler_discrete"] = ["FlowMapEulerDiscreteScheduler"] _import_structure["scheduling_flow_match_euler_discrete"] = ["FlowMatchEulerDiscreteScheduler"] _import_structure["scheduling_flow_match_heun_discrete"] = ["FlowMatchHeunDiscreteScheduler"] @@ -144,6 +145,7 @@ from .scheduling_entropy_bound import EntropyBoundScheduler, EntropyBoundSchedulerOutput from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler from .scheduling_euler_discrete import EulerDiscreteScheduler + from .scheduling_euler_lite import EulerLiteScheduler from .scheduling_flow_map_euler_discrete import FlowMapEulerDiscreteScheduler from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from .scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler diff --git a/src/diffusers/schedulers/scheduling_euler_lite.py b/src/diffusers/schedulers/scheduling_euler_lite.py new file mode 100644 index 000000000000..512b021ec4e5 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_euler_lite.py @@ -0,0 +1,137 @@ +# Copyright 2025 Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple, Union + +import torch + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin, SchedulerOutput + + +class EulerLiteScheduler(SchedulerMixin, ConfigMixin): + """First-order Euler ODE sampler over a linear-beta sigma schedule. + + Lite subset of [`EulerDiscreteScheduler`]: epsilon prediction, leading timesteps, and a + terminal sigma of 0. The update is the Karras Euler step from + `scheduling_euler_discrete.py` (no ancestral noise, no Karras/exponential/beta sigma + conversions). + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass + documentation for the generic methods the library implements for all schedulers such as + loading and saving. + + Args: + num_train_timesteps (`int`, defaults to `1000`): + The number of diffusion steps to train the model. + beta_start (`float`, defaults to `0.0001`): + The starting `beta` value of the linear schedule. + beta_end (`float`, defaults to `0.02`): + The final `beta` value of the linear schedule. + """ + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.0001, + beta_end: float = 0.02, + ): + self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) + self.alphas = 1.0 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) + # Same sigma definition as EulerDiscreteScheduler: sqrt((1 - bar_alpha) / bar_alpha). + self.sigmas_train = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 + self.timesteps = torch.arange(num_train_timesteps - 1, -1, -1) + self.sigmas = torch.cat([self.sigmas_train.flip(0), torch.zeros(1)]) + self.num_inference_steps: Optional[int] = None + + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device, None] = None): + """Set the discrete timesteps used for the denoising loop. + + Leading spacing matches `EulerDiscreteScheduler` with `timestep_spacing="leading"`: + `t = (arange(num_inference_steps) * (num_train_timesteps // num_inference_steps))` + reversed. Sigmas are the training-schedule values at those indices plus a trailing 0. + + Args: + num_inference_steps: Number of diffusion steps used at inference. + device: Device the timesteps should be moved to. + """ + self.num_inference_steps = num_inference_steps + step_ratio = self.config.num_train_timesteps // num_inference_steps + timesteps = (torch.arange(0, num_inference_steps) * step_ratio).round().flip(0).to(torch.long) + self.timesteps = timesteps.to(device) if device is not None else timesteps + sigmas = self.sigmas_train[self.timesteps.cpu()] + self.sigmas = torch.cat([sigmas, torch.zeros(1, dtype=sigmas.dtype)]) + + def step( + self, + model_output: torch.Tensor, + timestep: int, + sample: torch.Tensor, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[SchedulerOutput, Tuple]: + """Predict the sample at the previous timestep with a first-order Euler ODE step. + + For epsilon prediction this is the update in `EulerDiscreteScheduler.step` with + `s_churn=0` (no stochastic churn, so `generator` is unused): + + x0 = sample - sigma * epsilon + d = (sample - x0) / sigma + prev_sample = sample + d * (sigma_next - sigma) + + Args: + model_output: Direct output from the learned diffusion model (epsilon). + timestep: The current discrete timestep in the diffusion chain. + sample: A current instance of a sample created by the diffusion process. + generator: A torch.Generator for reproducible sampling. The Euler ODE step is + deterministic; this argument exists to match the scheduler `step` contract. + return_dict: Whether to return a SchedulerOutput or a plain tuple. + + Returns: + SchedulerOutput or tuple with the predicted previous sample. + """ + schedule = self.timesteps + if isinstance(timestep, torch.Tensor): + t_val = timestep.to(device=schedule.device, dtype=schedule.dtype).reshape(-1)[0] + else: + t_val = torch.tensor(timestep, device=schedule.device, dtype=schedule.dtype) + + matches = (schedule == t_val).nonzero() + if len(matches) > 0: + step_index = int(matches[0].item()) + sigma = self.sigmas[step_index] + sigma_next = self.sigmas[step_index + 1] + else: + t_int = int(t_val.item()) + sigma = self.sigmas_train[t_int] + if t_int > 0: + sigma_next = self.sigmas_train[t_int - 1] + else: + sigma_next = torch.zeros((), dtype=sigma.dtype, device=sigma.device) + + # Upcast to avoid precision issues when computing prev_sample (EulerDiscreteScheduler). + sample = sample.to(torch.float32) + sigma = sigma.to(device=sample.device, dtype=sample.dtype) + sigma_next = sigma_next.to(device=sample.device, dtype=sample.dtype) + + pred_original_sample = sample - sigma * model_output.to(sample.dtype) + derivative = (sample - pred_original_sample) / sigma + prev_sample = sample + derivative * (sigma_next - sigma) + prev_sample = prev_sample.to(model_output.dtype) + + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 3434c6416cce..11e55c9758a5 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -3471,6 +3471,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class EulerLiteScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class FlowMapEulerDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/schedulers/test_scheduling_euler_lite.py b/tests/schedulers/test_scheduling_euler_lite.py new file mode 100644 index 000000000000..96d2f1a33d35 --- /dev/null +++ b/tests/schedulers/test_scheduling_euler_lite.py @@ -0,0 +1,136 @@ +"""TEMPLATE — copy to tests/schedulers/test_scheduling_.py and set the +two constants below. Runs with zero installs (structural + signature contract) +and adds behavioral checks when torch + diffusers are available. + +See examples/scaffolded_scheduler + tests/schedulers/test_scheduling_ddpm_lite.py +for a filled-in example. +""" +import ast +import importlib.util +import unittest +from pathlib import Path + +# ---- edit these two for your scheduler ------------------------------------- +TARGET = Path(__file__).resolve().parents[2] / "src" / "diffusers" / "schedulers" / "scheduling_euler_lite.py" +CLASS = "EulerLiteScheduler" +# ---------------------------------------------------------------------------- + + +def _class_node(): + for node in ast.walk(ast.parse(TARGET.read_text())): + if isinstance(node, ast.ClassDef) and node.name == CLASS: + return node + raise AssertionError(f"{CLASS} not found in {TARGET}") + + +def _method(node, name): + for n in node.body: + if isinstance(n, ast.FunctionDef) and n.name == name: + return n + return None + + +class TestStructuralContract(unittest.TestCase): + """Thin structural sanity — the convention gate is the primary owner.""" + + @classmethod + def setUpClass(cls): + cls.node = _class_node() + + def test_inherits_required_mixins(self): + bases = {getattr(b, "id", getattr(b, "attr", "")) for b in self.node.bases} + self.assertEqual({"SchedulerMixin", "ConfigMixin"} & bases, {"SchedulerMixin", "ConfigMixin"}) + + def test_public_export(self): + root = Path(__file__).resolve().parents[2] + schedulers_init = (root / "src" / "diffusers" / "schedulers" / "__init__.py").read_text() + package_init = (root / "src" / "diffusers" / "__init__.py").read_text() + self.assertIn("scheduling_euler_lite", schedulers_init) + self.assertIn("EulerLiteScheduler", schedulers_init) + self.assertIn("EulerLiteScheduler", package_init) + + +class TestSignatureContract(unittest.TestCase): + """Deeper than the gate: verify the SIGNATURES, not just method presence.""" + + @classmethod + def setUpClass(cls): + cls.node = _class_node() + + def test_set_timesteps_signature(self): + m = _method(self.node, "set_timesteps") + self.assertIsNotNone(m, "set_timesteps missing") + args = [a.arg for a in m.args.args] + self.assertIn("num_inference_steps", args, + "set_timesteps must accept num_inference_steps (verified against upstream source)") + self.assertIn("device", args, "set_timesteps must accept device") + + def test_step_signature(self): + m = _method(self.node, "step") + self.assertIsNotNone(m, "step missing") + args = [a.arg for a in m.args.args] + for expected in ("model_output", "timestep", "sample", "generator"): + self.assertIn(expected, args, f"step must accept {expected} (upstream contract)") + + +class TestBehavioralContract(unittest.TestCase): + """Real behaviour — runs in CI with torch; skipped cleanly offline.""" + + def setUp(self): + if importlib.util.find_spec("torch") is None: + self.skipTest("torch not installed; behavioral contract runs in CI") + if importlib.util.find_spec("diffusers") is None: + self.skipTest("diffusers not installed; behavioral contract runs in CI") + spec = importlib.util.spec_from_file_location(CLASS, TARGET) + self.mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(self.mod) + self.Scheduler = getattr(self.mod, CLASS) + + def test_config_roundtrip(self): + s = self.Scheduler(num_train_timesteps=500) + rebuilt = self.Scheduler.from_config(s.config) + self.assertEqual(rebuilt.config.num_train_timesteps, 500) + + def test_timesteps_count(self): + s = self.Scheduler() + s.set_timesteps(25) + self.assertEqual(len(s.timesteps), 25) + + def test_output_type(self): + import torch + s = self.Scheduler() + s.set_timesteps(10) + sample = torch.zeros(1, 3, 8, 8) + out = s.step(torch.ones_like(sample), 1, sample, generator=torch.Generator().manual_seed(0)) + self.assertTrue(hasattr(out, "prev_sample")) + self.assertEqual(out.prev_sample.shape, sample.shape) + self.assertEqual(out.prev_sample.dtype, sample.dtype) + + def test_same_seed_same_output(self): + import torch + s = self.Scheduler() + s.set_timesteps(10) + sample = torch.zeros(1, 3, 8, 8) + mo = torch.ones_like(sample) + a = s.step(mo, 1, sample, generator=torch.Generator().manual_seed(0)).prev_sample + b = s.step(mo, 1, sample, generator=torch.Generator().manual_seed(0)).prev_sample + self.assertTrue(torch.equal(a, b)) + + def test_euler_ode_epsilon_update(self): + import torch + s = self.Scheduler() + s.set_timesteps(10) + sample = torch.zeros(1, 3, 8, 8) + model_output = torch.ones_like(sample) + t = int(s.timesteps[0].item()) + out = s.step(model_output, t, sample, generator=torch.Generator().manual_seed(0)).prev_sample + sigma = s.sigmas[0].to(device=sample.device, dtype=torch.float32) + sigma_next = s.sigmas[1].to(device=sample.device, dtype=torch.float32) + expected = sample.to(torch.float32) + (sigma_next - sigma) * model_output.to(torch.float32) + self.assertEqual(out.shape, sample.shape) + self.assertEqual(out.dtype, sample.dtype) + self.assertTrue(torch.allclose(out, expected.to(out.dtype))) + + +if __name__ == "__main__": + unittest.main()