Skip to content

Conversation

@Beinsezii
Copy link
Owner

@Beinsezii Beinsezii commented Sep 21, 2025

  • Functional sampling structures
  • Move all of the non-functional samplers
  • Runge-Kutta sampler
    • Better substep interpolation than 1d linear?
    • Unified derivative
  • Adaptive sampling?
    • Heun-Euler might be easy.
      • Just kind of yolo'd the evaluation. Might have to check what other impls do.
    • Should refactor tableau processing now that they're used twice.
  • Self-accurate unit tests
    • Diffusers unit tests? FlowMatchHeunScheduler might work at least.
  • Plots
  • Examples
  • Investigate running in Diffusers
    • Possibly could replace the LoopSequentialPipelineBlocks portions of ModularPipeline with a generic wrapper since they're small
      • A demo was created at examples/diffusers/functional.py using Flux. The API is very rough so I will probably hold off on making this an official interface until diffusers figures out their modular system more. Maybe I should open an issue suggesting a more functional-compatible system?
      • I opened a discussion upstream Functional Sampling for Modular Diffusers huggingface/diffusers#12405 to see if this is even the right path.
    • For samplers with a pre-determined step size ie Runge-Kutta I guess I could try the native diffusers higher order code by having SkrampleWrapperScheduler present some kind of generator over its own step function as model(x, t)?
      • This assumes they handle timesteps as an arbitrary sequence and don't assume it's whole steps which would only work for Heun. It would also involve pre-building the whole schedule splattered across the tableau, including skips for T=0.
      • It would also need to have another callback for all intermediate steps which is crazy unfortunate. I would honestly rather intercept the Euler calls and just copy the intermediate terms out of them.

@Beinsezii Beinsezii mentioned this pull request Sep 21, 2025
@Beinsezii Beinsezii mentioned this pull request Nov 5, 2025
Done
- Rename SkrampleSampler -> StructuredSampler in structured module
- Add FunctionalSampler and functional module
- Add Heun sampler in functional module
- Add StructuredFunctionalAdapter in interface module
- Add UT for StructuredFunctionalAdapter
- Add example for StructuredFunctionalAdapter and Heun
- Move some types and functions around to dedupe code.

Not Done
- Plots for functional samplers
- UT for Heun if diffusers behaves
- Higher order Heun / Runge Kutta
  - Ideally RKX instead of just RK3 or RK4
- Functional integrations?
- In *theory* it's possible to use FunctionalSinglestep through
  SkrampleWrapperScheduler by maintaining directly using `step` and
  doing a shitload of sketchy state management. Image2Image will 100%
  not work OOTB unless the pipeline actually respects `scheduler.order`
  which not all of them do I don't think. Alternatives involve monkey
  patches and sadness so I'd rather not.
Relies on branching for transforms currently, but results are
substantially better than the raw prediction based solver
The polar to/from derive functions are *so close* to working with flow
match but it's just not quite there.

So I basically have two options:
 1. fix the derives so they work for any transform
 2. accept defeat and replace SigmaTransform with some generic
    ModelTransforms unit struct that just has static methods
    for the sigma and derivative transforms
Beinsezii and others added 15 commits November 17, 2025 01:14
"Once I add velocity this can trivially go into the diffusers wrapper"

yeah pretty sure it's actually impossible to write v-prediction as

z̃ₛ = z̃ₜ + output · (ηₛ - ηₜ)

so I rewrote it again into a custom parameterization space that actually
for real works with everything and am wiping out all of the commits
trying to make DiffusionFlow work

Squashed commit of the following:

commit 6b83bf297e87b888a65733e1ed0c95b30f21bbd0
Author: Beinsezii <[email protected]>
Date:   Mon Nov 17 01:11:55 2025 -0800

    Rewrite ModelTransform to use custom parameterization

commit 4913744
Author: Beinsezii <[email protected]>
Date:   Sun Nov 16 16:45:17 2025 -0800

    Velocity yes depend on X

commit 4373aa1
Author: Beinsezii <[email protected]>
Date:   Fri Nov 14 22:56:11 2025 -0800

    Velocity don't depend on Epsilon

commit 0d0a4c0
Author: Beinsezii <[email protected]>
Date:   Fri Nov 14 22:22:55 2025 -0800

    Fix FlowModel to_x/from_x with VariancePreserving schedule

commit d6b5971
Author: Beinsezii <[email protected]>
Date:   Fri Nov 14 03:49:50 2025 -0800

    Make DiffusionModel the X impl

commit 73a73d6
Author: Beinsezii <[email protected]>
Date:   Fri Nov 14 03:22:07 2025 -0800

    models: Add XModel, VelocityModel, .to_x(), .from_x(), .to_h()

commit 26b012a
Author: Beinsezii <[email protected]>
Date:   Tue Nov 11 01:38:37 2025 -0800

    Add unified set of Diffusion model transforms

    Currently just for sampling.functional

    Once I add velocity this can trivially go into the diffusers wrapper
    instead of a raw prediction function, and then moved through structured

    Possibly I could associate the schedule and therefore also the sigma
    transform `type base_schedule = Linear` or whatever. Will have to play
    around see if FlowModel can be made to handle a Variance-Preserving
    schedule and vice-versa.

    Just so fucking glad I don't have to switch on sigma_transform anymore
    for the tableau solver.
Still not 100% certain on how to best parameterize the derivative
transform. It needs to be per-step for plans on making RKUltra its own
diffusers wrapper, so ig adding yet another parameter™ is hot.
Since we have common.divf() now this is more appropriate
Add unified set of Diffusion model transforms electric boogaloo
@Beinsezii Beinsezii marked this pull request as ready for review November 30, 2025 02:10
@Beinsezii
Copy link
Owner Author

Think I've added pretty much everything I could need for a first implementation. RKMoire and FastHeun are both crazy jank still but not sure what to do about that since diffusion is really chaotic.

@Beinsezii
Copy link
Owner Author

Beinsezii commented Nov 30, 2025

Basically, a lot of areas need follow-ups but most of it is no longer problems with the functional module but rather other areas of the library that need improved robustness as a result of findings in this PR:

  • ModelTransform invasion
    • Should be sent through to StructuredSampler and anywhere else that currently expects an X prediction
    • Will get super hairy with UniPC and DPM to avoid double-converts
  • Paramtetrize all tests Parametrize Tests #58
    • Much easier to debug large refactors. Should start with this
    • Possibly also works better with the parallel extension.
  • Possibly make schedules continuous over T
    • Mostly depends if Scaled can be made a continuous function.
      • Would probably need to hugely boost the test epsilons because I'd need to remove the round()
    • Would allow sending SkrampleScheduler directly instead of FloatSchedule, basically eliminating the sigma_transform variable. Would no longer be any point in assigning schedule into FunctionalSampler.
  • Possibly make RNG continuous over T
    • Might be better to do this in numpy/scipy?
    • Or I could actually make BrownianInterval the base class that Pyramid/Offset inherit from.
  • Possibly add RKUltraWrapper
    • By taking tableaux without overlapping timesteps like RK3.Ralston, it could be done purely with exact T indexing as we currently do.
  • Possibly add ComfyUI nodes
    • Requested
  • Add more obscure tableaux as I find them
    • Saw references to Feagin embedded 14/12 in some papers which sounds baller as hell but I really need a nice premade tableaux to have any hope of making it correct.
  • Improve RKMoire usability
    • Saw some junk in papers that could help but I'll need QwenVL to explain it to me as a child tbh.
  • Add DPM++ Singlestep methods
    • Mostly interesting because it implements stochastic sampling. Otherwise I'm not sure it's better than RKUltra.
  • Add whatever this is Res samplers? #52

@Beinsezii Beinsezii merged commit 40c5de0 into master Nov 30, 2025
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants