mirror of https://github.com/vladmandic/automatic
update schedulers and lint everything
Signed-off-by: vladmandic <mandic00@live.com>pull/4619/head
vladmandic
parent
f439d51ea7
commit
d7ca4f63a7
11
CHANGELOG.md
11
CHANGELOG.md
|
|
@ -1,6 +1,15 @@
|
|||
# Change Log for SD.Next
|
||||
|
||||
## Update for 2026-02-02
|
||||
## Update for 2026-02-04
|
||||
|
||||
### Highlights for 2026-02-04
|
||||
|
||||
Refresh release two weeks after prior release, yet we still somehow managed to pack in ~140 commits with new features, models and fixes!
|
||||
For full list of changes, see full changelog.
|
||||
|
||||
[ReadMe](https://github.com/vladmandic/automatic/blob/master/README.md) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) | [Sponsor](https://github.com/sponsors/vladmandic)
|
||||
|
||||
### Details for 2026-02-04
|
||||
|
||||
- **Models**
|
||||
- [Tongyi-MAI Z-Image Base](https://tongyi-mai.github.io/Z-Image-blog/)
|
||||
|
|
|
|||
167
TODO.md
167
TODO.md
|
|
@ -1,96 +1,123 @@
|
|||
# TODO
|
||||
|
||||
## Project Board
|
||||
|
||||
- <https://github.com/users/vladmandic/projects>
|
||||
|
||||
## Internal
|
||||
|
||||
- Feature: Flow-match `res4lyf` schedulers
|
||||
- Feature: Move `nunchaku` models to refernce instead of internal decision
|
||||
Project board: <https://github.com/users/vladmandic/projects>
|
||||
|
||||
- Update: `transformers==5.0.0`
|
||||
- Feature: Unify *huggingface* and *diffusers* model folders
|
||||
- Reimplement `llama` remover for Kanvas
|
||||
- Deploy: Create executable for SD.Next
|
||||
- Feature: Integrate natural language image search
|
||||
[ImageDB](https://github.com/vladmandic/imagedb)
|
||||
- Feature: Remote Text-Encoder support
|
||||
- Refactor: move sampler options to settings to config
|
||||
- Refactor: [GGUF](https://huggingface.co/docs/diffusers/main/en/quantization/gguf)
|
||||
- Feature: LoRA add OMI format support for SD35/FLUX.1
|
||||
- Refactor: remove `CodeFormer`
|
||||
- Refactor: remove `GFPGAN`
|
||||
- UI: Lite vs Expert mode
|
||||
- Video tab: add full API support
|
||||
- Control tab: add overrides handling
|
||||
- Engine: `TensorRT` acceleration
|
||||
- Deploy: Lite vs Expert mode
|
||||
- Engine: [mmgp](https://github.com/deepbeepmeep/mmgp)
|
||||
- Engine: [sharpfin](https://github.com/drhead/sharpfin) instead of `torchvision`
|
||||
- Engine: `TensorRT` acceleration
|
||||
- Feature: Auto handle scheduler `prediction_type`
|
||||
- Feature: Cache models in memory
|
||||
- Feature: Control tab: add overrides handling
|
||||
- Feature: Integrate natural language image search
|
||||
[ImageDB](https://github.com/vladmandic/imagedb)
|
||||
- Feature: LoRA add OMI format support for SD35/FLUX.1
|
||||
- Feature: Multi-user support
|
||||
- Feature: Remote Text-Encoder support
|
||||
- Feature: Settings profile manager
|
||||
- Feature: Video tab: add full API support
|
||||
- Refactor: Unify *huggingface* and *diffusers* model folders
|
||||
- Refactor: Move `nunchaku` models to refernce instead of internal decision
|
||||
- Refactor: [GGUF](https://huggingface.co/docs/diffusers/main/en/quantization/gguf)
|
||||
- Refactor: move sampler options to settings to config
|
||||
- Refactor: remove `CodeFormer`
|
||||
- Refactor: remove `GFPGAN`
|
||||
- Reimplement `llama` remover for Kanvas
|
||||
|
||||
## Modular
|
||||
|
||||
- Switch to modular pipelines
|
||||
- Feature: Transformers unified cache handler
|
||||
- Refactor: [Modular pipelines and guiders](https://github.com/huggingface/diffusers/issues/11915)
|
||||
- [MagCache](https://github.com/lllyasviel/FramePack/pull/673/files)
|
||||
- [SmoothCache](https://github.com/huggingface/diffusers/issues/11135)
|
||||
|
||||
## Features
|
||||
|
||||
- [Flux.2 TinyVAE](https://huggingface.co/fal/FLUX.2-Tiny-AutoEncoder)
|
||||
- [IPAdapter composition](https://huggingface.co/ostris/ip-composition-adapter)
|
||||
- [IPAdapter negative guidance](https://github.com/huggingface/diffusers/discussions/7167)
|
||||
- [STG](https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance)
|
||||
- [Video Inpaint Pipeline](https://github.com/huggingface/diffusers/pull/12506)
|
||||
- [Sonic Inpaint](https://github.com/ubc-vision/sonic)
|
||||
|
||||
### New models / Pipelines
|
||||
## New models / Pipelines
|
||||
|
||||
TODO: Investigate which models are diffusers-compatible and prioritize!
|
||||
|
||||
- [Bria FiboEdit](https://github.com/huggingface/diffusers/commit/d7a1c31f4f85bae5a9e01cdce49bd7346bd8ccd6)
|
||||
- [LTXVideo 0.98 LongMulti](https://github.com/huggingface/diffusers/pull/12614)
|
||||
- [Cosmos-Predict-2.5](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B)
|
||||
- [NewBie Image Exp0.1](https://github.com/huggingface/diffusers/pull/12803)
|
||||
- [Sana-I2V](https://github.com/huggingface/diffusers/pull/12634#issuecomment-3540534268)
|
||||
### Text-to-Image
|
||||
- [Bria FIBO](https://huggingface.co/briaai/FIBO)
|
||||
- [Bytedance Lynx](https://github.com/bytedance/lynx)
|
||||
- [ByteDance OneReward](https://github.com/bytedance/OneReward)
|
||||
- [ByteDance USO](https://github.com/bytedance/USO)
|
||||
- [Chroma Radiance](https://huggingface.co/lodestones/Chroma1-Radiance)
|
||||
- [Chroma Zeta](https://huggingface.co/lodestones/Zeta-Chroma)
|
||||
- [DiffSynth Studio](https://github.com/modelscope/DiffSynth-Studio)
|
||||
- [DiffusionForcing](https://github.com/kwsong0113/diffusion-forcing-transformer)
|
||||
- [Dream0 guidance](https://huggingface.co/ByteDance/DreamO)
|
||||
- [HunyuanAvatar](https://huggingface.co/tencent/HunyuanVideo-Avatar)
|
||||
- [HunyuanCustom](https://github.com/Tencent-Hunyuan/HunyuanCustom)
|
||||
- [Inf-DiT](https://github.com/zai-org/Inf-DiT)
|
||||
- [Krea Realtime Video](https://huggingface.co/krea/krea-realtime-video)
|
||||
- [LanDiff](https://github.com/landiff/landiff)
|
||||
- [Liquid](https://github.com/FoundationVision/Liquid)
|
||||
- [LongCat-Video](https://huggingface.co/meituan-longcat/LongCat-Video)
|
||||
- [LucyEdit](https://github.com/huggingface/diffusers/pull/12340)
|
||||
- [Chroma Radiance](https://huggingface.co/lodestones/Chroma1-Radiance)
|
||||
- [Liquid](https://github.com/FoundationVision/Liquid)
|
||||
- [Lumina-DiMOO](https://huggingface.co/Alpha-VLLM/Lumina-DiMOO)
|
||||
- [Magi](https://github.com/SandAI-org/MAGI-1)(https://github.com/huggingface/diffusers/pull/11713)
|
||||
- [Ming](https://github.com/inclusionAI/Ming)
|
||||
|
||||
### Image-to-Image
|
||||
- [Meituan LongCat-Image-Edit-Turbo](https://huggingface.co/meituan-longcat/LongCat-Image-Edit-Turbo)
|
||||
- [VIBE Image-Edit](https://huggingface.co/iitolstykh/VIBE-Image-Edit)
|
||||
- [Bria FiboEdit (diffusers)](https://github.com/huggingface/diffusers/commit/d7a1c31f4f85bae5a9e01cdce49bd7346bd8ccd6)
|
||||
- [LucyEdit (diffusers PR)](https://github.com/huggingface/diffusers/pull/12340)
|
||||
- [SD3 UltraEdit](https://github.com/HaozheZhao/UltraEdit)
|
||||
- [Step1X-Edit](https://github.com/stepfun-ai/Step1X-Edit)
|
||||
- [OneReward (mask-guided / object removal)](https://github.com/bytedance/OneReward)
|
||||
|
||||
### Text-to-Video
|
||||
- [OpenMOSS MOVA](https://huggingface.co/OpenMOSS-Team/MOVA-720p)
|
||||
- [Wan family (Wan2.1 / Wan2.2 variants)](https://huggingface.co/Wan-AI/Wan2.2-Animate-14B)
|
||||
- example: [Wan2.1-T2V-14B-CausVid](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-CausVid)
|
||||
- distill / step-distill examples: [Wan2.1-StepDistill-CfgDistill](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-StepDistill-CfgDistill)
|
||||
- [Krea Realtime Video](https://huggingface.co/krea/krea-realtime-video)
|
||||
- [MAGI-1 (autoregressive video)](https://github.com/SandAI-org/MAGI-1)
|
||||
- [MUG-V 10B (video generation)](https://huggingface.co/MUG-V/MUG-V-inference)
|
||||
- [Ovi (audio/video generation)](https://github.com/character-ai/Ovi)
|
||||
|
||||
### Image-to-Video
|
||||
- [MUG-V 10B](https://huggingface.co/MUG-V/MUG-V-inference)
|
||||
- [Ovi](https://github.com/character-ai/Ovi)
|
||||
- [Phantom HuMo](https://github.com/Phantom-video/Phantom)
|
||||
- [SD3 UltraEdit](https://github.com/HaozheZhao/UltraEdit)
|
||||
- [SelfForcing](https://github.com/guandeh17/Self-Forcing)
|
||||
- [SEVA](https://github.com/huggingface/diffusers/pull/11440)
|
||||
- [Step1X](https://github.com/stepfun-ai/Step1X-Edit)
|
||||
- [Wan-2.2 Animate](https://github.com/huggingface/diffusers/pull/12526)
|
||||
- [Wan-2.2 S2V](https://github.com/huggingface/diffusers/pull/12258)
|
||||
- [WAN-CausVid-Plus t2v](https://github.com/goatWu/CausVid-Plus/)
|
||||
- [WAN-CausVid](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-CausVid)
|
||||
- [WAN-StepDistill](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-StepDistill-CfgDistill)
|
||||
- [Wan2.2-Animate-14B](https://huggingface.co/Wan-AI/Wan2.2-Animate-14B)
|
||||
- [WAN2GP](https://github.com/deepbeepmeep/Wan2GP)
|
||||
- [HunyuanVideo-Avatar / HunyuanCustom](https://huggingface.co/tencent/HunyuanVideo-Avatar)
|
||||
- [Sana Image→Video (Sana-I2V)](https://github.com/huggingface/diffusers/pull/12634#issuecomment-3540534268)
|
||||
- [Wan-2.2 S2V (diffusers PR)](https://github.com/huggingface/diffusers/pull/12258)
|
||||
|
||||
### Migration
|
||||
### Video Editing / Long-Video / Animation Tooling
|
||||
- [LongCat-Video](https://huggingface.co/meituan-longcat/LongCat-Video)
|
||||
- [LTXVideo / LTXVideo LongMulti (diffusers PR)](https://github.com/huggingface/diffusers/pull/12614)
|
||||
- [DiffSynth-Studio (ModelScope)](https://github.com/modelscope/DiffSynth-Studio)
|
||||
- [Phantom (Phantom HuMo)](https://github.com/Phantom-video/Phantom)
|
||||
- [CausVid-Plus / WAN-CausVid-Plus](https://github.com/goatWu/CausVid-Plus/)
|
||||
- [Wan2GP (workflow/GUI for Wan)](https://github.com/deepbeepmeep/Wan2GP)
|
||||
|
||||
#### Asyncio
|
||||
### Multimodal
|
||||
- [Cosmos-Predict-2.5 (NVIDIA)](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B)
|
||||
- [Liquid (unified multimodal generator)](https://github.com/FoundationVision/Liquid)
|
||||
- [Lumina-DiMOO](https://huggingface.co/Alpha-VLLM/Lumina-DiMOO)
|
||||
- [Ming (inclusionAI)](https://github.com/inclusionAI/Ming)
|
||||
- [Magi (SandAI)](https://github.com/SandAI-org/MAGI-1)
|
||||
- [DreamO (ByteDance) — image customization framework](https://huggingface.co/ByteDance/DreamO)
|
||||
|
||||
### Other/Unsorted
|
||||
- [DiffusionForcing](https://github.com/kwsong0113/diffusion-forcing-transformer)
|
||||
- [Self-Forcing](https://github.com/guandeh17/Self-Forcing)
|
||||
- [SEVA](https://github.com/huggingface/diffusers/pull/11440)
|
||||
- [ByteDance USO](https://github.com/bytedance/USO)
|
||||
- [ByteDance Lynx](https://github.com/bytedance/lynx)
|
||||
- [LanDiff](https://github.com/landiff/landiff)
|
||||
- [MagCache](https://github.com/huggingface/diffusers/pull/12744)
|
||||
- [SmoothCache](https://github.com/huggingface/diffusers/issues/11135)
|
||||
- [STG](https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance)
|
||||
- [Video Inpaint Pipeline](https://github.com/huggingface/diffusers/pull/12506)
|
||||
- [Sonic Inpaint](https://github.com/ubc-vision/sonic)
|
||||
- [BoxDiff](https://github.com/huggingface/diffusers/pull/7947)
|
||||
- [Make-It-Count](https://github.com/Litalby1/make-it-count)
|
||||
- [FreeCustom](https://github.com/aim-uofa/FreeCustom)
|
||||
- [ControlNeXt](https://github.com/dvlab-research/ControlNeXt/)
|
||||
- [MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion)
|
||||
- [UniRef](https://github.com/FoundationVision/UniRef)
|
||||
- [AnyDoor](https://github.com/ali-vilab/AnyDoor)
|
||||
- [AnyText2](https://github.com/tyxsspa/AnyText2)
|
||||
- [DragonDiffusion](https://github.com/MC-E/DragonDiffusion)
|
||||
- [DenseDiffusion](https://github.com/naver-ai/DenseDiffusion)
|
||||
- [FlashFace](https://github.com/ali-vilab/FlashFace)
|
||||
- [PowerPaint](https://github.com/open-mmlab/PowerPaint)
|
||||
- [IC-Light](https://github.com/lllyasviel/IC-Light)
|
||||
- [ReNO](https://github.com/ExplainableML/ReNO)
|
||||
- [LoRAdapter](https://github.com/CompVis/LoRAdapter)
|
||||
- [LivePortrait](https://github.com/KwaiVGI/LivePortrait)
|
||||
|
||||
## Migration
|
||||
|
||||
### Asyncio
|
||||
|
||||
- Policy system is deprecated and will be removed in **Python 3.16**
|
||||
- [Python 3.14 removals - asyncio](https://docs.python.org/3.14/whatsnew/3.14.html#id10)
|
||||
|
|
@ -102,7 +129,7 @@ TODO: Investigate which models are diffusers-compatible and prioritize!
|
|||
- [asyncio.run](https://docs.python.org/3.14/library/asyncio-runner.html#asyncio.run)
|
||||
- [asyncio.Runner](https://docs.python.org/3.14/library/asyncio-runner.html#asyncio.Runner)
|
||||
|
||||
#### rmtree
|
||||
### rmtree
|
||||
|
||||
- `onerror` deprecated and replaced with `onexc` in **Python 3.12**
|
||||
``` python
|
||||
|
|
|
|||
|
|
@ -665,7 +665,7 @@ def check_diffusers():
|
|||
t_start = time.time()
|
||||
if args.skip_all:
|
||||
return
|
||||
sha = 'd4f97d19210d4abc32dac78fe7080cd8f5f0809c' # diffusers commit hash
|
||||
sha = '430c557b6a66a3c2b5740fb186324cb8a9f0f2e9' # diffusers commit hash
|
||||
# if args.use_rocm or args.use_zluda or args.use_directml:
|
||||
# sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now
|
||||
pkg = pkg_resources.working_set.by_key.get('diffusers', None)
|
||||
|
|
@ -811,6 +811,7 @@ def install_rocm_zluda():
|
|||
torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://rocm.nightlies.amd.com/{device.therock}')
|
||||
else:
|
||||
check_python(supported_minors=[12], reason='ROCm: Windows preview python==3.12 required')
|
||||
# torch 2.8.0a0 is the last version with rocm 6.4 support
|
||||
torch_command = os.environ.get('TORCH_COMMAND', '--no-cache-dir https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torch-2.8.0a0%2Bgitfc14c65-cp312-cp312-win_amd64.whl https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torchvision-0.24.0a0%2Bc85f008-cp312-cp312-win_amd64.whl')
|
||||
else:
|
||||
#check_python(supported_minors=[10, 11, 12, 13, 14], reason='ROCm backend requires a Python version between 3.10 and 3.13')
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from typing import List
|
|||
|
||||
|
||||
def xyz_grid_enum(option: str = "") -> List[dict]:
|
||||
from scripts.xyz import xyz_grid_classes
|
||||
from scripts.xyz import xyz_grid_classes # pylint: disable=no-name-in-module
|
||||
options = []
|
||||
for x in xyz_grid_classes.axis_options:
|
||||
_option = {
|
||||
|
|
|
|||
|
|
@ -1,24 +0,0 @@
|
|||
# TASK: Schedulers
|
||||
|
||||
## Notes
|
||||
|
||||
This is a codebase for diffusion schedulers implemented for `diffusers` library and ported from `res4lyf` repository at <https://github.com/ClownsharkBatwing/RES4LYF>
|
||||
|
||||
Ported schedulers codebase is in `modules/res4lyf`, do not modify any other files
|
||||
|
||||
## Testing
|
||||
|
||||
All schedulers were tested using prediction type `epsilon` and `StableDiffusionXLPipeline` pipeline for *text2image*: WORKING GOOD!
|
||||
|
||||
Shifting focus to testing prediction type `flow_prediction` and `ZImagePipeline` pipeline for *text2image*
|
||||
|
||||
## Results
|
||||
|
||||
- so far all tested schedules produce blocky/pixelated and unresolved output
|
||||
|
||||
## TODO
|
||||
|
||||
- [x] focus on a single scheduler only. lets pick abnorsett_2m (Fixed: Implemented AB update branch)
|
||||
- [x] validate config params: is this ok? (Validated: Config is correct for Flux/SD3 with new patch)
|
||||
config={'num_train_timesteps': 1000, 'beta_start': 0.0001, 'beta_end': 0.02, 'beta_schedule': 'linear', 'prediction_type': 'flow_prediction', 'variant': 'abnorsett_2m', 'use_analytic_solution': True, 'timestep_spacing': 'linspace', 'steps_offset': 0, 'use_flow_sigmas': True, 'shift': 3, 'base_shift': 0.5, 'max_shift': 1.15, 'base_image_seq_len': 256, 'max_image_seq_len': 4096}
|
||||
- [x] check code (Complete)
|
||||
|
|
@ -3,7 +3,7 @@
|
|||
from .abnorsett_scheduler import ABNorsettScheduler
|
||||
from .bong_tangent_scheduler import BongTangentScheduler
|
||||
from .common_sigma_scheduler import CommonSigmaScheduler
|
||||
from .deis_scheduler_alt import DEISMultistepScheduler
|
||||
from .deis_scheduler_alt import RESDEISMultistepScheduler
|
||||
from .etdrk_scheduler import ETDRKScheduler
|
||||
from .gauss_legendre_scheduler import GaussLegendreScheduler
|
||||
from .langevin_dynamics_scheduler import LangevinDynamicsScheduler
|
||||
|
|
@ -95,7 +95,7 @@ __all__ = [ # noqa: RUF022
|
|||
"RESMultistepSDEScheduler",
|
||||
"RESSinglestepScheduler",
|
||||
"RESSinglestepSDEScheduler",
|
||||
"DEISMultistepScheduler",
|
||||
"RESDEISMultistepScheduler",
|
||||
"ETDRKScheduler",
|
||||
"LawsonScheduler",
|
||||
"ABNorsettScheduler",
|
||||
|
|
@ -183,7 +183,7 @@ BASE = [
|
|||
("RES Multistep SDE", RESMultistepSDEScheduler),
|
||||
("RES Singlestep", RESSinglestepScheduler),
|
||||
("RES Singlestep SDE", RESSinglestepSDEScheduler),
|
||||
("DEIS Multistep", DEISMultistepScheduler),
|
||||
("DEIS Multistep", RESDEISMultistepScheduler),
|
||||
("ETDRK", ETDRKScheduler),
|
||||
("Lawson", LawsonScheduler),
|
||||
("ABNorsett", ABNorsettScheduler),
|
||||
|
|
|
|||
|
|
@ -123,7 +123,6 @@ class ABNorsettScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing {self.config.timestep_spacing} is not supported.")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
|
||||
|
||||
if self.config.use_karras_sigmas:
|
||||
|
|
@ -135,8 +134,10 @@ class ABNorsettScheduler(SchedulerMixin, ConfigMixin):
|
|||
elif self.config.use_flow_sigmas:
|
||||
s_min = getattr(self.config, "sigma_min", None)
|
||||
s_max = getattr(self.config, "sigma_max", None)
|
||||
if s_min is None: s_min = 0.001
|
||||
if s_max is None: s_max = 1.0
|
||||
if s_min is None:
|
||||
s_min = 0.001
|
||||
if s_max is None:
|
||||
s_max = 1.0
|
||||
sigmas = np.linspace(s_max, s_min, num_inference_steps)
|
||||
|
||||
if self.config.shift != 1.0 or self.config.use_dynamic_shifting:
|
||||
|
|
@ -279,11 +280,11 @@ class ABNorsettScheduler(SchedulerMixin, ConfigMixin):
|
|||
# AB2 Variable Step
|
||||
# x_{n+1} = x_n + dt * [ (1 + r/2) * v_n - (r/2) * v_{n-1} ]
|
||||
# where r = dt_cur / dt_prev
|
||||
|
||||
|
||||
v_nm1 = self.model_outputs[-2]
|
||||
sigma_prev = self.prev_sigmas[-2]
|
||||
dt_prev = sigma_curr - sigma_prev
|
||||
|
||||
|
||||
if abs(dt_prev) < 1e-8:
|
||||
# Fallback to Euler if division by zero risk
|
||||
x_next = sample + dt * v_n
|
||||
|
|
@ -293,15 +294,15 @@ class ABNorsettScheduler(SchedulerMixin, ConfigMixin):
|
|||
c0 = 1 + 0.5 * r
|
||||
c1 = -0.5 * r
|
||||
x_next = sample + dt * (c0 * v_n + c1 * v_nm1)
|
||||
|
||||
|
||||
elif curr_order >= 3:
|
||||
# For now, fallback to AB2 (variable) for higher orders to ensure stability
|
||||
# given the complexity of variable-step AB3/4 formulas inline.
|
||||
# For now, fallback to AB2 (variable) for higher orders to ensure stability
|
||||
# given the complexity of variable-step AB3/4 formulas inline.
|
||||
# The user specifically requested abnorsett_2m.
|
||||
v_nm1 = self.model_outputs[-2]
|
||||
sigma_prev = self.prev_sigmas[-2]
|
||||
dt_prev = sigma_curr - sigma_prev
|
||||
|
||||
|
||||
if abs(dt_prev) < 1e-8:
|
||||
x_next = sample + dt * v_n
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -99,6 +99,8 @@ class BongTangentScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
@ -154,8 +156,8 @@ class BongTangentScheduler(SchedulerMixin, ConfigMixin):
|
|||
end_val = sigma_min
|
||||
mid_val = sigma_mid
|
||||
|
||||
tan_sigmas_1 = self._get_bong_tangent_sigmas(stage_1_len, s1, p1, start_val, mid_val)
|
||||
tan_sigmas_2 = self._get_bong_tangent_sigmas(stage_2_len, s2, p2 - stage_1_len, mid_val, end_val)
|
||||
tan_sigmas_1 = self._get_bong_tangent_sigmas(stage_1_len, s1, p1, start_val, mid_val, dtype=dtype)
|
||||
tan_sigmas_2 = self._get_bong_tangent_sigmas(stage_2_len, s2, p2 - stage_1_len, mid_val, end_val, dtype=dtype)
|
||||
|
||||
tan_sigmas_1 = tan_sigmas_1[:-1]
|
||||
sigmas_list = tan_sigmas_1 + tan_sigmas_2
|
||||
|
|
@ -208,7 +210,7 @@ class BongTangentScheduler(SchedulerMixin, ConfigMixin):
|
|||
from .scheduler_utils import add_noise_to_sample
|
||||
return add_noise_to_sample(original_samples, noise, self.sigmas, timesteps, self.timesteps)
|
||||
|
||||
def _get_bong_tangent_sigmas(self, steps: int, slope: float, pivot: int, start: float, end: float) -> List[float]:
|
||||
def _get_bong_tangent_sigmas(self, steps: int, slope: float, pivot: int, start: float, end: float, dtype: torch.dtype = torch.float32) -> List[float]:
|
||||
x = torch.arange(steps, dtype=dtype)
|
||||
|
||||
def bong_fn(val):
|
||||
|
|
|
|||
|
|
@ -203,6 +203,8 @@ class CommonSigmaScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ def get_def_integral_3(a, b, c, start, end, d):
|
|||
|
||||
class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
||||
"""
|
||||
DEISMultistepScheduler: Diffusion Explicit Iterative Sampler with high-order multistep.
|
||||
RESDEISMultistepScheduler: Diffusion Explicit Iterative Sampler with high-order multistep.
|
||||
Adapted from the RES4LYF repository.
|
||||
"""
|
||||
|
||||
|
|
@ -82,6 +82,7 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
self._step_index = None
|
||||
self._sigmas_cpu = None
|
||||
self.all_coeffs = []
|
||||
self.prev_sigmas = []
|
||||
|
||||
def set_timesteps(
|
||||
self,
|
||||
|
|
@ -109,7 +110,7 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
timesteps = np.maximum(timesteps, 0)
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
log_sigmas_all = np.log(np.maximum(sigmas, 1e-10))
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -244,7 +245,7 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
step_index = self._step_index
|
||||
sigma_t = self.sigmas[step_index]
|
||||
|
||||
|
||||
# RECONSTRUCT X0 (Matching PEC pattern)
|
||||
if self.config.prediction_type == "epsilon":
|
||||
denoised = sample - sigma_t * model_output
|
||||
|
|
@ -262,13 +263,10 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
if self.config.clip_sample:
|
||||
denoised = denoised.clamp(-self.config.sample_max_value, self.config.sample_max_value)
|
||||
|
||||
# DEIS coefficients are precomputed in set_timesteps
|
||||
coeffs = self.all_coeffs[step_index]
|
||||
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
# Variable Step Adams-Bashforth for Flow Matching
|
||||
self.model_outputs.append(model_output)
|
||||
self.prev_sigmas.append(sigma_t)
|
||||
self.prev_sigmas.append(sigma_t)
|
||||
# Note: deis uses hist_samples for x0? I'll use model_outputs for v.
|
||||
if len(self.model_outputs) > 4:
|
||||
self.model_outputs.pop(0)
|
||||
|
|
@ -276,7 +274,7 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
dt = self.sigmas[step_index + 1] - sigma_t
|
||||
v_n = model_output
|
||||
|
||||
|
||||
curr_order = min(len(self.prev_sigmas), 3)
|
||||
|
||||
if curr_order == 1:
|
||||
|
|
@ -301,11 +299,11 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
x_next = sample + dt * (c0 * v_n + c1 * self.model_outputs[-2])
|
||||
|
||||
self._step_index += 1
|
||||
if not return_dict: return (x_next,)
|
||||
if not return_dict:
|
||||
return (x_next,)
|
||||
return SchedulerOutput(prev_sample=x_next)
|
||||
|
||||
sigma_next = self.sigmas[step_index + 1]
|
||||
alpha_next = 1 / (sigma_next**2 + 1) ** 0.5 if sigma_next > 0 else 1.0
|
||||
|
||||
if self.config.solver_order == 1:
|
||||
# 1st order step (Euler) in x-space
|
||||
|
|
@ -316,11 +314,11 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
h = -torch.log(sigma_next / sigma_t) if sigma_t > 0 and sigma_next > 0 else torch.zeros_like(sigma_t)
|
||||
phi = Phi(h, [0], getattr(self.config, "use_analytic_solution", True))
|
||||
phi_1 = phi(1)
|
||||
|
||||
|
||||
# History of denoised samples
|
||||
x0s = [denoised] + self.model_outputs[::-1]
|
||||
orders = min(len(x0s), self.config.solver_order)
|
||||
|
||||
|
||||
# Force Order 1 at the end of schedule
|
||||
if self.num_inference_steps is not None and step_index >= self.num_inference_steps - 3:
|
||||
res = phi_1 * denoised
|
||||
|
|
@ -334,7 +332,7 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
h_prev = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 1] + 1e-9))
|
||||
h_prev_t = torch.tensor(h_prev, device=sample.device, dtype=sample.dtype)
|
||||
r = h_prev_t / (h + 1e-9)
|
||||
|
||||
|
||||
# Hard Restart
|
||||
if r < 0.5 or r > 2.0:
|
||||
res = phi_1 * denoised
|
||||
|
|
@ -355,7 +353,7 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
h_p2 = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 2] + 1e-9))
|
||||
r1 = torch.tensor(h_p1, device=sample.device, dtype=sample.dtype) / (h + 1e-9)
|
||||
r2 = torch.tensor(h_p2, device=sample.device, dtype=sample.dtype) / (h + 1e-9)
|
||||
|
||||
|
||||
# Hard Restart
|
||||
if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0:
|
||||
res = phi_1 * denoised
|
||||
|
|
|
|||
|
|
@ -174,6 +174,8 @@ class ETDRKScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -165,7 +165,6 @@ class GaussLegendreScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing must be one of 'linspace', 'leading', or 'trailing', got {self.config.timestep_spacing}")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -323,7 +322,7 @@ class GaussLegendreScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -190,6 +190,8 @@ class LangevinDynamicsScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -172,6 +172,8 @@ class LawsonScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -121,7 +121,6 @@ class LinearRKScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing must be one of 'linspace', 'leading', or 'trailing', got {self.config.timestep_spacing}")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -262,7 +261,7 @@ class LinearRKScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -121,7 +121,6 @@ class LobattoScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing must be one of 'linspace', 'leading', or 'trailing', got {self.config.timestep_spacing}")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -262,7 +261,7 @@ class LobattoScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -180,6 +180,8 @@ class PECScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
return sample / ((sigma**2 + 1) ** 0.5)
|
||||
|
||||
|
|
|
|||
|
|
@ -155,7 +155,6 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing must be one of 'linspace', 'leading', or 'trailing', got {self.config.timestep_spacing}")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -247,11 +246,7 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin):
|
|||
sigma = self.sigmas[self._step_index]
|
||||
return sample / ((sigma**2 + 1) ** 0.5)
|
||||
|
||||
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
||||
from .scheduler_utils import index_for_timestep
|
||||
if schedule_timesteps is None:
|
||||
schedule_timesteps = self.timesteps
|
||||
return index_for_timestep(timestep, schedule_timesteps)
|
||||
|
||||
|
||||
def _init_step_index(self, timestep):
|
||||
if self._step_index is None:
|
||||
|
|
@ -309,7 +304,7 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -250,7 +250,7 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
# Variable Step Adams-Bashforth for Flow Matching
|
||||
dt = sigma_next - sigma
|
||||
v_n = model_output
|
||||
|
||||
|
||||
if curr_order == 1:
|
||||
x_next = sample + dt * v_n
|
||||
elif curr_order == 2:
|
||||
|
|
@ -258,7 +258,7 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
sigma_prev = self.prev_sigmas[-2]
|
||||
dt_prev = sigma - sigma_prev
|
||||
r = dt / dt_prev if abs(dt_prev) > 1e-8 else 0.0
|
||||
|
||||
|
||||
# Stability check
|
||||
if dt_prev == 0 or r < -0.9 or r > 2.0: # Fallback
|
||||
x_next = sample + dt * v_n
|
||||
|
|
@ -267,17 +267,6 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
c1 = -0.5 * r
|
||||
x_next = sample + dt * (c0 * v_n + c1 * self.model_outputs[-2])
|
||||
elif curr_order >= 3:
|
||||
# AB3
|
||||
sigma_prev1 = self.prev_sigmas[-2]
|
||||
sigma_prev2 = self.prev_sigmas[-3]
|
||||
dt_prev1 = sigma - sigma_prev1
|
||||
dt_prev2 = self.prev_sigmas[-2] - sigma_prev2 # This is not strictly correct for variable steps logic used in ABNorsett, assume simplified AB3 for now or stick to AB2
|
||||
# Actually, let's reuse ABNorsett logic
|
||||
# x_{n+1} = x_n + dt * [ (1 + r1/2 + r2/2 + ... ) ] - Too complex to derive on the fly?
|
||||
# Let's use AB2 for stability as requested "Variable Step Adams-Bashforth like ABNorsett"
|
||||
# ABNorsett implemented AB2. I will downgrade order 3 to AB2 for safety or implement AB3 if confident.
|
||||
# Let's stick to AB2 for Flow as it is robust enough.
|
||||
|
||||
# Re-implement AB2 logic
|
||||
sigma_prev = self.prev_sigmas[-2]
|
||||
dt_prev = sigma - sigma_prev
|
||||
|
|
@ -291,7 +280,7 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
self.model_outputs.pop(0)
|
||||
self.x0_outputs.pop(0)
|
||||
self.prev_sigmas.pop(0)
|
||||
|
||||
|
||||
if not return_dict:
|
||||
return (x_next,)
|
||||
return SchedulerOutput(prev_sample=x_next)
|
||||
|
|
@ -301,25 +290,16 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
phi_1 = phi(1)
|
||||
|
||||
if variant.startswith("res"):
|
||||
# REiS Multistep logic
|
||||
c2, c3 = 0.5, 1.0
|
||||
|
||||
# Force Order 1 at the end of schedule
|
||||
if self.num_inference_steps is not None and self._step_index >= self.num_inference_steps - 3:
|
||||
curr_order = 1
|
||||
|
||||
if curr_order == 2:
|
||||
h_prev = -torch.log(self.prev_sigmas[-1] / self.prev_sigmas[-2])
|
||||
c2 = (-h_prev / h).item() if h > 0 else 0.5
|
||||
rk_type = "res_2s"
|
||||
elif curr_order == 3:
|
||||
h_prev1 = -torch.log(self.prev_sigmas[-1] / self.prev_sigmas[-2])
|
||||
h_prev2 = -torch.log(self.prev_sigmas[-1] / self.prev_sigmas[-3])
|
||||
c2 = (-h_prev1 / h).item() if h > 0 else 0.5
|
||||
c3 = (-h_prev2 / h).item() if h > 0 else 1.0
|
||||
rk_type = "res_3s"
|
||||
pass
|
||||
else:
|
||||
rk_type = "res_1s"
|
||||
pass
|
||||
|
||||
# Exponential Integrator Update in x-space
|
||||
if curr_order == 1:
|
||||
|
|
@ -329,21 +309,21 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
# b2 = -phi_2 / r = -phi(2) / (h_prev/h)
|
||||
# Here we use: b2 = phi(2) / ((-h_prev / h) + 1e-9)
|
||||
# Since (-h_prev/h) is negative (-r), this gives correct negative sign for b2.
|
||||
|
||||
|
||||
# Stability check
|
||||
r_check = h_prev / (h + 1e-9) # This is effectively -r if using h_prev definition above?
|
||||
# Wait, h_prev above is -log(). Positive.
|
||||
# h is positive.
|
||||
# So h_prev/h is positive. defined as r in other files.
|
||||
# But here code uses -h_prev / h in denominator.
|
||||
|
||||
|
||||
# Stability check
|
||||
r_check = h_prev / (h + 1e-9)
|
||||
|
||||
|
||||
# Hard Restart
|
||||
if r_check < 0.5 or r_check > 2.0:
|
||||
res = phi_1 * x0
|
||||
else:
|
||||
else:
|
||||
b2 = phi(2) / ((-h_prev / h) + 1e-9)
|
||||
b1 = phi_1 - b2
|
||||
res = b1 * self.x0_outputs[-1] + b2 * self.x0_outputs[-2]
|
||||
|
|
@ -353,7 +333,7 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
h_p2 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-3] + 1e-9))
|
||||
r1 = h_p1 / (h + 1e-9)
|
||||
r2 = h_p2 / (h + 1e-9)
|
||||
|
||||
|
||||
if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0:
|
||||
res = phi_1 * x0
|
||||
else:
|
||||
|
|
@ -429,13 +409,14 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
elif order == 2:
|
||||
h_prev = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9))
|
||||
r = h_prev / (h + 1e-9)
|
||||
|
||||
|
||||
# Correct Adams-Bashforth-like coefficients for Exponential Integrators
|
||||
|
||||
|
||||
# Hard Restart for stability
|
||||
if r < 0.5 or r > 2.0:
|
||||
return [[phi_1]]
|
||||
|
||||
phi_2 = phi(2)
|
||||
b2 = -phi_2 / (r + 1e-9)
|
||||
b1 = phi_1 - b2
|
||||
return [[b1, b2]]
|
||||
|
|
@ -444,13 +425,13 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin):
|
|||
h_prev2 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-3] + 1e-9))
|
||||
r1 = h_prev1 / (h + 1e-9)
|
||||
r2 = h_prev2 / (h + 1e-9)
|
||||
|
||||
|
||||
if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0:
|
||||
return [[phi_1]]
|
||||
|
||||
|
||||
phi_2 = phi(2)
|
||||
phi_3 = phi(3)
|
||||
|
||||
|
||||
# Generalized AB3 for Exponential Integrators (Varying steps)
|
||||
denom = r2 - r1 + 1e-9
|
||||
b3 = (phi_3 + r1 * phi_2) / (r2 * denom)
|
||||
|
|
|
|||
|
|
@ -105,6 +105,8 @@ class RESMultistepSDEScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -89,8 +89,6 @@ class RESSinglestepScheduler(SchedulerMixin, ConfigMixin):
|
|||
self._begin_index = begin_index
|
||||
|
||||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
|
|
@ -124,12 +122,11 @@ class RESSinglestepScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing {self.config.timestep_spacing} is not supported.")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
|
||||
# Linear remapping logic
|
||||
if self.config.use_flow_sigmas:
|
||||
# Logic handled below (linspace) or here?
|
||||
# To match others:
|
||||
pass
|
||||
# Logic handled here
|
||||
pass
|
||||
else:
|
||||
sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
|
||||
|
||||
|
|
@ -152,10 +149,11 @@ class RESSinglestepScheduler(SchedulerMixin, ConfigMixin):
|
|||
self.config.base_image_seq_len,
|
||||
self.config.max_image_seq_len,
|
||||
)
|
||||
if not self.config.use_flow_sigmas:
|
||||
sigmas = apply_shift(torch.from_numpy(sigmas), shift).numpy()
|
||||
|
||||
if self.config.use_flow_sigmas:
|
||||
sigmas = np.linspace(1.0, 1 / 1000, num_inference_steps)
|
||||
else:
|
||||
sigmas = apply_shift(torch.from_numpy(sigmas), shift).numpy()
|
||||
|
||||
if self.config.use_flow_sigmas:
|
||||
timesteps = sigmas * self.config.num_train_timesteps
|
||||
|
|
@ -211,12 +209,13 @@ class RESSinglestepScheduler(SchedulerMixin, ConfigMixin):
|
|||
x0 = sample - sigma * model_output
|
||||
else:
|
||||
x0 = model_output
|
||||
|
||||
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
dt = sigma_next - sigma
|
||||
x_next = sample + dt * model_output
|
||||
self._step_index += 1
|
||||
if not return_dict: return (x_next,)
|
||||
if not return_dict:
|
||||
return (x_next,)
|
||||
return SchedulerOutput(prev_sample=x_next)
|
||||
|
||||
# Exponential Integrator Update
|
||||
|
|
|
|||
|
|
@ -93,6 +93,8 @@ class RESSinglestepSDEScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -185,7 +185,7 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin):
|
|||
# phi_2 = phi(2) # Moved inside conditional blocks as needed
|
||||
|
||||
history_len = len(self.x0_outputs)
|
||||
|
||||
|
||||
# Stability: Force Order 1 for final few steps to prevent degradation at low noise levels
|
||||
if self.num_inference_steps is not None and self._step_index >= self.num_inference_steps - 3:
|
||||
return [phi_1], h
|
||||
|
|
@ -193,10 +193,10 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin):
|
|||
if self.config.rk_type in ["res_2m", "deis_2m"] and history_len >= 2:
|
||||
h_prev = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9))
|
||||
r = h_prev / (h + 1e-9)
|
||||
|
||||
|
||||
h_prev = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9))
|
||||
r = h_prev / (h + 1e-9)
|
||||
|
||||
|
||||
# Hard Restart: if step sizes vary too wildly, fallback to order 1
|
||||
if r < 0.5 or r > 2.0:
|
||||
return [phi_1], h
|
||||
|
|
@ -220,10 +220,10 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin):
|
|||
# Hard Restart check
|
||||
if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0:
|
||||
return [phi_1], h
|
||||
|
||||
|
||||
phi_2 = phi(2)
|
||||
phi_3 = phi(3)
|
||||
|
||||
|
||||
# Generalized AB3 for Exponential Integrators (Varying steps)
|
||||
denom = r2 - r1 + 1e-9
|
||||
b3 = (phi_3 + r1 * phi_2) / (r2 * denom)
|
||||
|
|
@ -275,9 +275,9 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin):
|
|||
# Variable Step Adams-Bashforth for Flow Matching
|
||||
dt = sigma_next - sigma
|
||||
v_n = model_output
|
||||
|
||||
|
||||
curr_order = min(len(self.prev_sigmas), 3) # Max order 3 here
|
||||
|
||||
|
||||
if curr_order == 1:
|
||||
x_next = sample + dt * v_n
|
||||
elif curr_order == 2:
|
||||
|
|
@ -300,7 +300,8 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin):
|
|||
x_next = sample + dt * (c0 * v_n + c1 * self.model_outputs[-2])
|
||||
|
||||
self._step_index += 1
|
||||
if not return_dict: return (x_next,)
|
||||
if not return_dict:
|
||||
return (x_next,)
|
||||
return SchedulerOutput(prev_sample=x_next)
|
||||
|
||||
# GET COEFFICIENTS
|
||||
|
|
|
|||
|
|
@ -205,6 +205,8 @@ class RiemannianFlowScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -75,7 +75,6 @@ class RungeKutta44Scheduler(SchedulerMixin, ConfigMixin):
|
|||
# 1. Base sigmas
|
||||
timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=float)[::-1].copy()
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -108,7 +107,6 @@ class RungeKutta44Scheduler(SchedulerMixin, ConfigMixin):
|
|||
sigmas_expanded.append(0.0) # terminal sigma
|
||||
|
||||
# 3. Map back to timesteps
|
||||
log_sigmas_all = np.log(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
sigmas_interpolated = np.array(sigmas_expanded)
|
||||
# Linear remapping for Flow Matching
|
||||
timesteps_expanded = sigmas_interpolated * self.config.num_train_timesteps
|
||||
|
|
@ -195,7 +193,7 @@ class RungeKutta44Scheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -94,7 +94,6 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin):
|
|||
timesteps = np.maximum(timesteps, 0)
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -145,7 +144,6 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin):
|
|||
sigmas_expanded.append(s_curr + c * (s_next - s_curr))
|
||||
sigmas_expanded.append(0.0)
|
||||
|
||||
log_sigmas_all = np.log(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
sigmas_interpolated = np.array(sigmas_expanded)
|
||||
# Linear remapping for Flow Matching
|
||||
timesteps_expanded = sigmas_interpolated * self.config.num_train_timesteps
|
||||
|
|
@ -198,6 +196,18 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin):
|
|||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
|
||||
# Dormand-Prince 5(4) Coefficients
|
||||
a = [
|
||||
[],
|
||||
[1/5],
|
||||
[3/40, 9/40],
|
||||
[44/45, -56/15, 32/9],
|
||||
[19372/6561, -25360/2187, 64448/6561, -212/729],
|
||||
[9017/3168, -355/33, 46732/5247, 49/176, -5103/18656],
|
||||
[35/384, 0, 500/1113, 125/192, -2187/6784, 11/84]
|
||||
]
|
||||
b = [35/384, 0, 500/1113, 125/192, -2187/6784, 11/84, 0]
|
||||
|
||||
step_index = self._step_index
|
||||
stage_index = step_index % 7
|
||||
|
||||
|
|
@ -229,7 +239,7 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -94,7 +94,6 @@ class RungeKutta67Scheduler(SchedulerMixin, ConfigMixin):
|
|||
timesteps = np.maximum(timesteps, 0)
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -145,7 +144,6 @@ class RungeKutta67Scheduler(SchedulerMixin, ConfigMixin):
|
|||
sigmas_expanded.append(s_curr + c * (s_next - s_curr))
|
||||
sigmas_expanded.append(0.0)
|
||||
|
||||
log_sigmas_all = np.log(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
sigmas_interpolated = np.array(sigmas_expanded)
|
||||
# Linear remapping for Flow Matching
|
||||
timesteps_expanded = sigmas_interpolated * self.config.num_train_timesteps
|
||||
|
|
@ -229,7 +227,7 @@ class RungeKutta67Scheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -90,21 +90,19 @@ def get_dynamic_shift(mu, base_shift, max_shift, base_seq_len, max_seq_len):
|
|||
return m * mu + b
|
||||
|
||||
def index_for_timestep(timestep, timesteps):
|
||||
import numpy as _np
|
||||
|
||||
# Normalize inputs to numpy arrays for a robust, device-agnostic argmin
|
||||
if isinstance(timestep, torch.Tensor):
|
||||
timestep_np = timestep.detach().cpu().numpy()
|
||||
else:
|
||||
timestep_np = _np.array(timestep)
|
||||
timestep_np = np.array(timestep)
|
||||
|
||||
if isinstance(timesteps, torch.Tensor):
|
||||
timesteps_np = timesteps.detach().cpu().numpy()
|
||||
else:
|
||||
timesteps_np = _np.array(timesteps)
|
||||
timesteps_np = np.array(timesteps)
|
||||
|
||||
# Use numpy argmin on absolute difference for stability
|
||||
idx = _np.abs(timesteps_np - timestep_np).argmin()
|
||||
idx = np.abs(timesteps_np - timestep_np).argmin()
|
||||
return int(idx)
|
||||
|
||||
def add_noise_to_sample(
|
||||
|
|
|
|||
|
|
@ -155,6 +155,8 @@ class SimpleExponentialScheduler(SchedulerMixin, ConfigMixin):
|
|||
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
|
||||
if self._step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
if self.config.prediction_type == "flow_prediction":
|
||||
return sample
|
||||
sigma = self.sigmas[self._step_index]
|
||||
sample = sample / ((sigma**2 + 1) ** 0.5)
|
||||
return sample
|
||||
|
|
|
|||
|
|
@ -125,7 +125,6 @@ class SpecializedRKScheduler(SchedulerMixin, ConfigMixin):
|
|||
raise ValueError(f"timestep_spacing must be one of 'linspace', 'leading', or 'trailing', got {self.config.timestep_spacing}")
|
||||
|
||||
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
|
||||
log_sigmas_all = np.log(sigmas)
|
||||
if self.config.interpolation_type == "linear":
|
||||
sigmas = np.interp(timesteps, np.arange(len(sigmas)), sigmas)
|
||||
elif self.config.interpolation_type == "log_linear":
|
||||
|
|
@ -286,7 +285,7 @@ class SpecializedRKScheduler(SchedulerMixin, ConfigMixin):
|
|||
|
||||
# derivative = (x - x0) / sigma
|
||||
derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample)
|
||||
|
||||
|
||||
if self.sample_at_start_of_step is None:
|
||||
if stage_index > 0:
|
||||
# Mid-step fallback for Img2Img/Inpainting
|
||||
|
|
|
|||
|
|
@ -0,0 +1,228 @@
|
|||
import os
|
||||
import sys
|
||||
import time
|
||||
import inspect
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Ensure we can import modules
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")))
|
||||
|
||||
from modules.errors import log
|
||||
from modules.res4lyf import (
|
||||
BASE, SIMPLE, VARIANTS,
|
||||
RESUnifiedScheduler, RESMultistepScheduler, RESDEISMultistepScheduler,
|
||||
ETDRKScheduler, LawsonScheduler, ABNorsettScheduler, PECScheduler,
|
||||
RiemannianFlowScheduler, RESSinglestepScheduler, RESSinglestepSDEScheduler,
|
||||
RESMultistepSDEScheduler, SimpleExponentialScheduler, LinearRKScheduler,
|
||||
LobattoScheduler, GaussLegendreScheduler, RungeKutta44Scheduler,
|
||||
RungeKutta57Scheduler, RungeKutta67Scheduler, SpecializedRKScheduler,
|
||||
BongTangentScheduler, CommonSigmaScheduler, RadauIIAScheduler,
|
||||
LangevinDynamicsScheduler
|
||||
)
|
||||
|
||||
def test_scheduler(name, scheduler_class, config):
|
||||
try:
|
||||
scheduler = scheduler_class(**config)
|
||||
except Exception as e:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} error="Init failed: {e}"')
|
||||
return False
|
||||
|
||||
num_steps = 20
|
||||
scheduler.set_timesteps(num_steps)
|
||||
|
||||
sample = torch.randn((1, 4, 64, 64))
|
||||
has_changed = False
|
||||
t0 = time.time()
|
||||
messages = []
|
||||
|
||||
try:
|
||||
for i, t in enumerate(scheduler.timesteps):
|
||||
# Simulate model output (noise or x0 or v), Using random noise for stability check
|
||||
model_output = torch.randn_like(sample)
|
||||
|
||||
# Scaling Check
|
||||
sigma = scheduler.sigmas[scheduler.step_index] if scheduler.step_index is not None else scheduler.sigmas[0] # Handle potential index mismatch if step_index is updated differently, usually step_index matches i for these tests
|
||||
|
||||
# Re-introduce scaling calculation first
|
||||
scaled_sample = scheduler.scale_model_input(sample, t)
|
||||
|
||||
if config.get("prediction_type") == "flow_prediction":
|
||||
expected_scale = 1.0
|
||||
else:
|
||||
expected_scale = 1.0 / ((sigma**2 + 1) ** 0.5)
|
||||
|
||||
# Simple check with loose tolerance due to float precision
|
||||
expected_scaled_sample = sample * expected_scale
|
||||
if not torch.allclose(scaled_sample, expected_scaled_sample, atol=1e-4):
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} expected={expected_scale} error="scaling mismatch"')
|
||||
return False
|
||||
|
||||
if torch.isnan(scaled_sample).any():
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="NaN in scaled_sample"')
|
||||
return False
|
||||
|
||||
if torch.isinf(scaled_sample).any():
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Inf in scaled_sample"')
|
||||
return False
|
||||
|
||||
output = scheduler.step(model_output, t, sample)
|
||||
|
||||
# Shape and Dtype check
|
||||
if output.prev_sample.shape != sample.shape:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Shape mismatch: {output.prev_sample.shape} vs {sample.shape}"')
|
||||
return False
|
||||
if output.prev_sample.dtype != sample.dtype:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Dtype mismatch: {output.prev_sample.dtype} vs {sample.dtype}"')
|
||||
return False
|
||||
|
||||
# Update check: Did the sample change?
|
||||
if not torch.equal(sample, output.prev_sample):
|
||||
has_changed = True
|
||||
|
||||
# Sample Evolution Check
|
||||
step_diff = (sample - output.prev_sample).abs().mean().item()
|
||||
if step_diff < 1e-6:
|
||||
messages.append(f'warning="minimal sample change: {step_diff}"')
|
||||
|
||||
sample = output.prev_sample
|
||||
|
||||
if torch.isnan(sample).any():
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="NaN in sample"')
|
||||
return False
|
||||
|
||||
if torch.isinf(sample).any():
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Inf in sample"')
|
||||
return False
|
||||
|
||||
# Divergence check
|
||||
if sample.abs().max() > 1e10:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="divergence detected"')
|
||||
return False
|
||||
|
||||
# External check for Sigma Monotonicity
|
||||
if hasattr(scheduler, 'sigmas'):
|
||||
sigmas = scheduler.sigmas.cpu().numpy()
|
||||
if len(sigmas) > 1:
|
||||
diffs = np.diff(sigmas) # Check if potentially monotonic decreasing (standard) OR increasing (some flow/inverse setups). We allow flat sections (diff=0) hence 1e-6 slack
|
||||
is_monotonic_decreasing = np.all(diffs <= 1e-6)
|
||||
is_monotonic_increasing = np.all(diffs >= -1e-6)
|
||||
if not (is_monotonic_decreasing or is_monotonic_increasing):
|
||||
messages.append('warning="sigmas are not monotonic"')
|
||||
|
||||
except Exception as e:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} exception: {e}')
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
if not has_changed:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} error="sample never changed"')
|
||||
return False
|
||||
|
||||
final_std = sample.std().item()
|
||||
with open("std_log.txt", "a") as f:
|
||||
f.write(f"STD_LOG: {name} config={config} std={final_std}\n")
|
||||
|
||||
if final_std > 50.0 or final_std < 0.1:
|
||||
log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} std={final_std} error="variance drift"')
|
||||
|
||||
t1 = time.time()
|
||||
messages = list(set(messages))
|
||||
log.info(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} time={t1-t0} messages={messages}')
|
||||
return True
|
||||
|
||||
def run_tests():
|
||||
prediction_types = ["epsilon", "v_prediction", "sample"] # flow_prediction is special, usually requires flow sigmas or specific setup, checking standard ones first
|
||||
|
||||
# Test BASE schedulers with their specific parameters
|
||||
log.warning('type="base"')
|
||||
for name, cls in BASE:
|
||||
configs = []
|
||||
|
||||
# prediction_types
|
||||
for pt in prediction_types:
|
||||
configs.append({"prediction_type": pt})
|
||||
|
||||
# Specific params for specific classes
|
||||
if cls == RESUnifiedScheduler:
|
||||
rk_types = ["res_2m", "res_3m", "res_2s", "res_3s", "res_5s", "res_6s", "deis_1s", "deis_2m", "deis_3m"]
|
||||
for rk in rk_types:
|
||||
for pt in prediction_types:
|
||||
configs.append({"rk_type": rk, "prediction_type": pt})
|
||||
|
||||
elif cls == RESMultistepScheduler:
|
||||
variants = ["res_2m", "res_3m", "deis_2m", "deis_3m"]
|
||||
for v in variants:
|
||||
for pt in prediction_types:
|
||||
configs.append({"variant": v, "prediction_type": pt})
|
||||
|
||||
elif cls == RESDEISMultistepScheduler:
|
||||
for order in range(1, 6):
|
||||
for pt in prediction_types:
|
||||
configs.append({"solver_order": order, "prediction_type": pt})
|
||||
|
||||
elif cls == ETDRKScheduler:
|
||||
variants = ["etdrk2_2s", "etdrk3_a_3s", "etdrk3_b_3s", "etdrk4_4s", "etdrk4_4s_alt"]
|
||||
for v in variants:
|
||||
for pt in prediction_types:
|
||||
configs.append({"variant": v, "prediction_type": pt})
|
||||
|
||||
elif cls == LawsonScheduler:
|
||||
variants = ["lawson2a_2s", "lawson2b_2s", "lawson4_4s"]
|
||||
for v in variants:
|
||||
for pt in prediction_types:
|
||||
configs.append({"variant": v, "prediction_type": pt})
|
||||
|
||||
elif cls == ABNorsettScheduler:
|
||||
variants = ["abnorsett_2m", "abnorsett_3m", "abnorsett_4m"]
|
||||
for v in variants:
|
||||
for pt in prediction_types:
|
||||
configs.append({"variant": v, "prediction_type": pt})
|
||||
|
||||
elif cls == PECScheduler:
|
||||
variants = ["pec423_2h2s", "pec433_2h3s"]
|
||||
for v in variants:
|
||||
for pt in prediction_types:
|
||||
configs.append({"variant": v, "prediction_type": pt})
|
||||
|
||||
elif cls == RiemannianFlowScheduler:
|
||||
metrics = ["euclidean", "hyperbolic", "spherical", "lorentzian"]
|
||||
for m in metrics:
|
||||
configs.append({"metric_type": m, "prediction_type": "epsilon"}) # Flow usually uses v or raw, but epsilon check matches others
|
||||
|
||||
if not configs:
|
||||
for pt in prediction_types:
|
||||
configs.append({"prediction_type": pt})
|
||||
|
||||
for conf in configs:
|
||||
test_scheduler(name, cls, conf)
|
||||
|
||||
log.warning('type="simple"')
|
||||
for name, cls in SIMPLE:
|
||||
for pt in prediction_types:
|
||||
test_scheduler(name, cls, {"prediction_type": pt})
|
||||
|
||||
log.warning('type="variants"')
|
||||
for name, cls in VARIANTS:
|
||||
# these classes preset their variants/rk_types in __init__ so we just test prediction types
|
||||
for pt in prediction_types:
|
||||
test_scheduler(name, cls, {"prediction_type": pt})
|
||||
|
||||
# Extra robustness check: Flow Prediction Type
|
||||
log.warning('type="flow"')
|
||||
flow_schedulers = [
|
||||
RESUnifiedScheduler, RESMultistepScheduler, ABNorsettScheduler,
|
||||
RESSinglestepScheduler, RESSinglestepSDEScheduler, RESDEISMultistepScheduler,
|
||||
RESMultistepSDEScheduler, ETDRKScheduler, LawsonScheduler, PECScheduler,
|
||||
SimpleExponentialScheduler, LinearRKScheduler, LobattoScheduler,
|
||||
GaussLegendreScheduler, RungeKutta44Scheduler, RungeKutta57Scheduler,
|
||||
RungeKutta67Scheduler, SpecializedRKScheduler, BongTangentScheduler,
|
||||
CommonSigmaScheduler, RadauIIAScheduler, LangevinDynamicsScheduler,
|
||||
RiemannianFlowScheduler
|
||||
]
|
||||
for cls in flow_schedulers:
|
||||
test_scheduler(cls.__name__, cls, {"prediction_type": "flow_prediction", "use_flow_sigmas": True})
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_tests()
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
from .abnorsett_scheduler import ABNorsettScheduler
|
||||
from .common_sigma_scheduler import CommonSigmaScheduler
|
||||
from .deis_scheduler_alt import DEISMultistepScheduler
|
||||
from .deis_scheduler_alt import RESDEISMultistepScheduler
|
||||
from .etdrk_scheduler import ETDRKScheduler
|
||||
from .gauss_legendre_scheduler import GaussLegendreScheduler
|
||||
from .lawson_scheduler import LawsonScheduler
|
||||
|
|
@ -302,19 +302,19 @@ class SigmaSmoothScheduler(CommonSigmaScheduler):
|
|||
super().__init__(**kwargs)
|
||||
|
||||
## DEIS Multistep Variants
|
||||
class DEIS1MultistepScheduler(DEISMultistepScheduler):
|
||||
class DEIS1MultistepScheduler(RESDEISMultistepScheduler):
|
||||
def __init__(self, **kwargs):
|
||||
kwargs["order"] = "1"
|
||||
kwargs["solver_order"] = 1
|
||||
super().__init__(**kwargs)
|
||||
|
||||
class DEIS2MultistepScheduler(DEISMultistepScheduler):
|
||||
class DEIS2MultistepScheduler(RESDEISMultistepScheduler):
|
||||
def __init__(self, **kwargs):
|
||||
kwargs["order"] = "2"
|
||||
kwargs["solver_order"] = 2
|
||||
super().__init__(**kwargs)
|
||||
|
||||
class DEIS3MultistepScheduler(DEISMultistepScheduler):
|
||||
class DEIS3MultistepScheduler(RESDEISMultistepScheduler):
|
||||
def __init__(self, **kwargs):
|
||||
kwargs["order"] = "3"
|
||||
kwargs["solver_order"] = 3
|
||||
super().__init__(**kwargs)
|
||||
|
||||
## Linear RK Variants
|
||||
|
|
|
|||
|
|
@ -59,7 +59,7 @@ try:
|
|||
from modules.schedulers.scheduler_ufogen import UFOGenScheduler # pylint: disable=ungrouped-imports
|
||||
from modules.schedulers.scheduler_unipc_flowmatch import FlowUniPCMultistepScheduler # pylint: disable=ungrouped-imports
|
||||
from modules.schedulers.scheduler_flashflow import FlashFlowMatchEulerDiscreteScheduler # pylint: disable=ungrouped-imports
|
||||
from modules.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports
|
||||
from modules.schedulers.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports
|
||||
except Exception as e:
|
||||
shared.log.error(f'Sampler import: version={diffusers.__version__} error: {e}')
|
||||
if os.environ.get('SD_SAMPLER_DEBUG', None) is not None:
|
||||
|
|
|
|||
|
|
@ -44,9 +44,9 @@ class Script(scripts_manager.Script):
|
|||
return None
|
||||
|
||||
import yaml
|
||||
from scripts.ctrlx import CtrlXStableDiffusionXLPipeline
|
||||
from scripts.ctrlx.sdxl import get_control_config, register_control
|
||||
from scripts.ctrlx.utils import get_self_recurrence_schedule
|
||||
from scripts.ctrlx import CtrlXStableDiffusionXLPipeline # pylint: disable=no-name-in-module
|
||||
from scripts.ctrlx.sdxl import get_control_config, register_control # pylint: disable=no-name-in-module
|
||||
from scripts.ctrlx.utils import get_self_recurrence_schedule # pylint: disable=no-name-in-module
|
||||
|
||||
orig_prompt_attention = shared.opts.prompt_attention
|
||||
shared.opts.data['prompt_attention'] = 'fixed'
|
||||
|
|
|
|||
|
|
@ -58,7 +58,7 @@ class Script(scripts_manager.Script):
|
|||
shared.log.warning('FreeScale: missing input image')
|
||||
return None
|
||||
|
||||
from scripts.freescale import StableDiffusionXLFreeScale, StableDiffusionXLFreeScaleImg2Img
|
||||
from scripts.freescale import StableDiffusionXLFreeScale, StableDiffusionXLFreeScaleImg2Img # pylint: disable=no-name-in-module
|
||||
self.orig_pipe = shared.sd_model
|
||||
self.orig_slice = shared.opts.diffusers_vae_slicing
|
||||
self.orig_tile = shared.opts.diffusers_vae_tiling
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ def verify_insightface():
|
|||
|
||||
|
||||
def load_infiniteyou(model: str):
|
||||
from scripts.infiniteyou import InfUFluxPipeline
|
||||
from scripts.infiniteyou import InfUFluxPipeline # pylint: disable=no-name-in-module
|
||||
shared.sd_model = InfUFluxPipeline(
|
||||
pipe=shared.sd_model,
|
||||
model_version=model,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from safetensors.torch import load_file
|
||||
from scripts.layerdiffuse.layerdiffuse_model import LoraLoader, AttentionSharingProcessor
|
||||
from scripts.layerdiffuse.layerdiffuse_model import LoraLoader, AttentionSharingProcessor # pylint: disable=no-name-in-module
|
||||
|
||||
|
||||
def merge_delta_weights_into_unet(pipe, delta_weights):
|
||||
|
|
|
|||
|
|
@ -65,7 +65,7 @@ class Script(scripts_manager.Script):
|
|||
if repo_id is not None:
|
||||
import huggingface_hub as hf
|
||||
repo_file = hf.snapshot_download(repo_id, cache_dir=shared.opts.hfcache_dir)
|
||||
from scripts.lbm import get_model
|
||||
from scripts.lbm import get_model # pylint: disable=no-name-in-module
|
||||
model = get_model(
|
||||
repo_file,
|
||||
save_dir=None,
|
||||
|
|
@ -85,7 +85,7 @@ class Script(scripts_manager.Script):
|
|||
install('lpips')
|
||||
|
||||
from torchvision.transforms import ToPILImage, ToTensor
|
||||
from scripts.lbm import get_model, extract_object, resize_and_center_crop
|
||||
from scripts.lbm import get_model, extract_object, resize_and_center_crop # pylint: disable=no-name-in-module
|
||||
|
||||
ori_h_bg, ori_w_bg = fg_image.size
|
||||
ar_bg = ori_h_bg / ori_w_bg
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ class Script(scripts_manager.Script):
|
|||
supported_model_list = ['sdxl']
|
||||
if shared.sd_model_type not in supported_model_list:
|
||||
shared.log.warning(f'PixelSmith: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}')
|
||||
from scripts.pixelsmith import PixelSmithXLPipeline, PixelSmithVAE
|
||||
from scripts.pixelsmith import PixelSmithXLPipeline, PixelSmithVAE # pylint: disable=no-name-in-module
|
||||
self.orig_pipe = shared.sd_model
|
||||
self.orig_vae = shared.sd_model.vae
|
||||
if self.vae is None:
|
||||
|
|
|
|||
|
|
@ -123,7 +123,7 @@ class StableDiffusionXLPuLIDPipeline:
|
|||
if sampler is not None:
|
||||
self.sampler = sampler
|
||||
else:
|
||||
from scripts.pulid import sampling
|
||||
from scripts.pulid import sampling # pylint: disable=no-name-in-module
|
||||
self.sampler = sampling.sample_dpmpp_sde
|
||||
|
||||
@property
|
||||
|
|
|
|||
|
|
@ -34,11 +34,11 @@ class Script(scripts_manager.Script):
|
|||
return model, sampler, width, height, start, scale, lora
|
||||
|
||||
def run(self, p: processing.StableDiffusionProcessing, model, sampler, width, height, start, scale, lora): # pylint: disable=arguments-differ, unused-argument
|
||||
from scripts.xadapter.xadapter_hijacks import PositionNet
|
||||
from scripts.xadapter.xadapter_hijacks import PositionNet # pylint: disable=no-name-in-module
|
||||
diffusers.models.embeddings.PositionNet = PositionNet # patch diffusers==0.26 from diffusers==0.20
|
||||
from scripts.xadapter.adapter import Adapter_XL
|
||||
from scripts.xadapter.pipeline_sd_xl_adapter import StableDiffusionXLAdapterPipeline
|
||||
from scripts.xadapter.unet_adapter import UNet2DConditionModel as UNet2DConditionModelAdapter
|
||||
from scripts.xadapter.adapter import Adapter_XL # pylint: disable=no-name-in-module
|
||||
from scripts.xadapter.pipeline_sd_xl_adapter import StableDiffusionXLAdapterPipeline # pylint: disable=no-name-in-module
|
||||
from scripts.xadapter.unet_adapter import UNet2DConditionModel as UNet2DConditionModelAdapter # pylint: disable=no-name-in-module
|
||||
|
||||
global adapter # pylint: disable=global-statement
|
||||
if model == 'None':
|
||||
|
|
|
|||
2
wiki
2
wiki
|
|
@ -1 +1 @@
|
|||
Subproject commit a678731c5da5d77f7c21edf7f1d62d8307d0d7fe
|
||||
Subproject commit da7620df144de8d2af259eff2b7a4522783f38cc
|
||||
Loading…
Reference in New Issue