Skip to content

Fix Cosmos2 safety checker grad mode leak#13799

Open
ooovenenoso wants to merge 1 commit into
huggingface:mainfrom
ooovenenoso:fix/cosmos2-grad-mode-safety-checker
Open

Fix Cosmos2 safety checker grad mode leak#13799
ooovenenoso wants to merge 1 commit into
huggingface:mainfrom
ooovenenoso:fix/cosmos2-grad-mode-safety-checker

Conversation

@ooovenenoso
Copy link
Copy Markdown

What does this PR do?

Fixes #13790.

Cosmos2VideoToWorldPipeline creates a default CosmosSafetyChecker when safety_checker=None. If the guardrail dependency changes PyTorch's global grad mode during construction, that state leaks out of from_pretrained() and leaves torch.is_grad_enabled() disabled globally.

This preserves and restores the caller's grad mode around default safety checker construction.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Was this discussed/approved via a Github issue? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

@DN6

Tests

  • Added test_default_safety_checker_preserves_grad_mode to cover default safety checker construction that disables grad mode.
  • Verified RED locally by running the new test before the fix and observing it fail with AssertionError: False is not true.
  • python -m pytest tests/pipelines/cosmos/test_cosmos2_video2world.py::Cosmos2VideoToWorldPipelineFastTests::test_default_safety_checker_preserves_grad_mode -q
  • python -m pytest tests/pipelines/cosmos/test_cosmos2_video2world.py -q
  • make style
  • make quality

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

loading Cosmos2 pipeline also disabled gradient tracking globally

1 participant