New model: Microsoft Lens

[dxqb]

Test in preview branch: https://github.com/Nerogar/OneTrainer/tree/preview

Summary

This PR adds Microsoft Lens: https://huggingface.co/microsoft/Lens

I'll leave this here as a draft for anyone who wants to experiment with this model, because it'll take a while until merge.

• Lens is not part of diffusers yet: [New Model] Add LensPipeline （microsoft/Lens） huggingface/diffusers#13837
• It uses Microsoft's upstream code instead: https://github.com/microsoft/Lens
• That currently requires transformers v5. Transformers v5 breaks other models: Upgrade transformers to 5.9 #1506

Open points:

• code review incomplete, there is still unchecked AI code
• Claude: Open point — timestep-shift schedule doesn't match the paper's training μ (§2.3)

Includes #1506 and #1509

Test plan

• pre-commit run --all-files passes
• Launched the affected UI or script and exercised the change
• Tested with at least one real preset / config when relevant (note which: Lens)

AI assistance

• Early AI prototype — opened for discussion, not ready for review

[Silvicultor]

Did a bit of testing with this PR:
What works:

• Basic LoRA training.
• Sampling.

What doesn’t work:

• Torch compile, unlike a similar bug with Anima PR this one is already triggered by compiling transformer blocks (not the optimizer). Attached the error log.
• Loading created LoRA weights in ComfyUI.

Loss starts on a pretty high level (~0.6). A longer run would be needed to see if that is a problem or not. But to do a longer run I would want common inference tools support the LoRA format first.

Traceback (most recent call last):
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1883, in _compile
    raise_unimplemented_cache_limit_exceeded()
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1867, in raise_unimplemented_cache_limit_exceeded
    unimplemented(
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/exc.py", line 657, in unimplemented
    raise Unsupported(msg, gb_type, skip_frame)
torch._dynamo.exc.Unsupported: Dynamo recompile limit exceeded
  Explanation: Dynamo attempted to recompile the code object too many times, exceeding the recompile_limit cache size limit (currently set to 8). Excessive recompilations can degrade performance due to the compilation overhead of each recompilation.
  Hint: To monitor recompilations, enable TORCH_LOGS=recompiles. If recompilations are expected, consider increasing torch._dynamo.config.recompile_limit to an appropriate value.
  Hint: See https://docs.pytorch.org/docs/main/user_guide/torch_compiler/compile/programming_model.recompilation.html for tips on dealing with recompilations.

  Developer debug context: Limit type: recompile_limit

 For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0039.html

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/user/AI/OT_Lens/OneTrainer/modules/ui/TrainUI.py", line 739, in __training_thread_function
    trainer.train()
  File "/home/user/AI/OT_Lens/OneTrainer/modules/trainer/GenericTrainer.py", line 756, in train
    loss.backward()
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_tensor.py", line 631, in backward
    torch.autograd.backward(
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/autograd/__init__.py", line 379, in backward
    _engine_run_backward(
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/autograd/graph.py", line 882, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/autograd/function.py", line 317, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/utils/checkpoint.py", line 314, in backward
    outputs = ctx.run_function(*detached_inputs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/modules/util/checkpointing_util.py", line 124, in __checkpointing_forward
    output = self.orig_forward(*args) if self.checkpoint is None else self.checkpoint(*args)
                                                                      ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1047, in compile_wrapper
    result = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 2474, in __call__
    result = self._torchdynamo_orig_backend(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 736, in __call__
    result = _compile(
             ^^^^^^^^^
  File "/home/user/AI/OT_Lens/OneTrainer/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1890, in _compile
    raise FailOnRecompileLimitHit(
torch._dynamo.exc.FailOnRecompileLimitHit: Hard failure due to fullgraph=True
Creating Backup workspace/run/backup/2026-06-07_11-12-09-backup-3-0-3
Saving models/model.safetensors

[dxqb]

Did a bit of testing with this PR

torch.compile: that's a torch bug, introduced in torch 2.12: pytorch/pytorch#186537
will need a workaround as long as we use torch 2.12. You could downgrade to torch 2.11 in the meantime.
Comfy: Lens is new - when I wrote this code, Comfy didn't even support the model at all yet. Will take care of that when the PR is close to merge.
high loss: that's normal for the Flux2 VAE, which Lens uses.

[dxqb]

torch.compile: that's a torch bug, introduced in torch 2.12: pytorch/pytorch#186537 will need a workaround as long as we use torch 2.12. You could downgrade to torch 2.11 in the meantime.

Merged the bug workaround branch into this branch #1511
I didn't see the error even before this workaround, so if you could re-test that would be useful. Quite sure it's fixed though.

[Silvicultor]

torch.compile: that's a torch bug, introduced in torch 2.12: pytorch/pytorch#186537 will need a workaround as long as we use torch 2.12. You could downgrade to torch 2.11 in the meantime.

Merged the bug workaround branch into this branch #1511 I didn't see the error even before this workaround, so if you could re-test that would be useful. Quite sure it's fixed though.

Yes, the workaround helps. Now both transformer blocks as well as the optimizer can be compiled.