Spaces:
Build error
Build error
File size: 2,214 Bytes
b6af722 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
import warnings
import attrs
from cosmos_predict1.utils import log
from cosmos_predict1.utils.config import CheckpointConfig as BaseCheckpointConfig
from cosmos_predict1.utils.config import make_freezable
from cosmos_predict1.utils.fsdp_checkpointer import FSDPCheckpointer as BaseFSDPCheckpointer
@make_freezable
@attrs.define(slots=False)
class CheckpointConfig(BaseCheckpointConfig):
load_ema_to_reg: bool = False
class FSDPCheckpointer(BaseFSDPCheckpointer):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if not isinstance(self.config_checkpoint, CheckpointConfig):
warnings.warn(
"The 'config_checkpoint' is not an instance of 'CheckpointConfig'. "
"This behavior is deprecated and will not be supported in future versions. "
"Please update 'config_checkpoint' to be of type 'CheckpointConfig'.",
DeprecationWarning,
)
self.load_ema_to_reg = False
else:
self.load_ema_to_reg = self.config_checkpoint.load_ema_to_reg
log.critical(f"load_ema_to_reg: {self.load_ema_to_reg}", rank0_only=False)
def load_model_during_init(self, model, is_ema: bool = False, ema_id: int = 0):
if self.load_ema_to_reg and is_ema is False:
is_ema = True
ema_id = 0
log.critical("Loading EMA model to regular model during initialization.", rank0_only=False)
super().load_model_during_init(model, is_ema, ema_id)
|