Spaces:

hugging-apps
/

fruc-collaborative-3d-reconstruction

Running on Zero

fruc-collaborative-3d-reconstruction / fruc /utils /checkpoint.py

Upload folder using huggingface_hub

3982600 verified 27 days ago

724 Bytes

"""Normalize the checkpoint containers used by the public FRUC entry points."""
from collections.abc import Mapping


def extract_state_dict(checkpoint):
    if not isinstance(checkpoint, Mapping):
        raise TypeError("Expected a state dictionary or a checkpoint mapping.")
    state = checkpoint
    for key in ("model_state_dict", "state_dict", "model"):
        if key in checkpoint and isinstance(checkpoint[key], Mapping):
            state = checkpoint[key]
            break
    if not state or any(not isinstance(key, str) for key in state):
        raise ValueError("Checkpoint does not contain a non-empty parameter dictionary.")
    return {key.removeprefix("module."): value for key, value in state.items()}

	"""Normalize the checkpoint containers used by the public FRUC entry points."""
	from collections.abc import Mapping


	def extract_state_dict(checkpoint):
	if not isinstance(checkpoint, Mapping):
	raise TypeError("Expected a state dictionary or a checkpoint mapping.")
	state = checkpoint
	for key in ("model_state_dict", "state_dict", "model"):
	if key in checkpoint and isinstance(checkpoint[key], Mapping):
	state = checkpoint[key]
	break
	if not state or any(not isinstance(key, str) for key in state):
	raise ValueError("Checkpoint does not contain a non-empty parameter dictionary.")
	return {key.removeprefix("module."): value for key, value in state.items()}