doanh25032004

Backup source tree of video_gen_physics (2026-07-31T14:21:08Z)

ec0a9aa verified 3 months ago

3.21 kB

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
Functions for downloading pre-trained DiT models
"""
import json
import os

import torch
from torchvision.datasets.utils import download_url

pretrained_models = {"DiT-XL-2-512x512.pt", "DiT-XL-2-256x256.pt"}


def find_model(model_name):
    """
    Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path.
    """
    if model_name in pretrained_models:  # Find/download our pre-trained DiT checkpoints
        return download_model(model_name)
    else:  # Load a custom DiT checkpoint:
        if not os.path.isfile(model_name):
            # if the model_name is a directory, then we assume we should load it in the Hugging Face manner
            # i.e. the model weights are sharded into multiple files and there is an index.json file
            # walk through the files in the directory and find the index.json file
            index_file = [os.path.join(model_name, f) for f in os.listdir(model_name) if "index.json" in f]
            assert len(index_file) == 1, f"Could not find index.json in {model_name}"

            # process index json
            with open(index_file[0], "r") as f:
                index_data = json.load(f)

            bin_to_weight_mapping = dict()
            for k, v in index_data["weight_map"].items():
                if v in bin_to_weight_mapping:
                    bin_to_weight_mapping[v].append(k)
                else:
                    bin_to_weight_mapping[v] = [k]

            # make state dict
            state_dict = dict()
            for bin_name, weight_list in bin_to_weight_mapping.items():
                bin_path = os.path.join(model_name, bin_name)
                bin_state_dict = torch.load(bin_path, map_location=lambda storage, loc: storage)
                for weight in weight_list:
                    state_dict[weight] = bin_state_dict[weight]
            return state_dict
        else:
            # if it is a file, we just load it directly in the typical PyTorch manner
            assert os.path.exists(model_name), f"Could not find DiT checkpoint at {model_name}"
            checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
            if "ema" in checkpoint:  # supports checkpoints from train.py
                checkpoint = checkpoint["ema"]
            return checkpoint


def download_model(model_name):
    """
    Downloads a pre-trained DiT model from the web.
    """
    assert model_name in pretrained_models
    local_path = f"pretrained_models/{model_name}"
    if not os.path.isfile(local_path):
        os.makedirs("pretrained_models", exist_ok=True)
        web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{model_name}"
        download_url(web_path, "pretrained_models")
    model = torch.load(local_path, map_location=lambda storage, loc: storage)
    return model


if __name__ == "__main__":
    # Download all DiT checkpoints
    for model in pretrained_models:
        download_model(model)
    print("Done.")

	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.

	# This source code is licensed under the license found in the
	# LICENSE file in the root directory of this source tree.

	"""
	Functions for downloading pre-trained DiT models
	"""
	import json
	import os

	import torch
	from torchvision.datasets.utils import download_url

	pretrained_models = {"DiT-XL-2-512x512.pt", "DiT-XL-2-256x256.pt"}


	def find_model(model_name):
	"""
	Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path.
	"""
	if model_name in pretrained_models: # Find/download our pre-trained DiT checkpoints
	return download_model(model_name)
	else: # Load a custom DiT checkpoint:
	if not os.path.isfile(model_name):
	# if the model_name is a directory, then we assume we should load it in the Hugging Face manner
	# i.e. the model weights are sharded into multiple files and there is an index.json file
	# walk through the files in the directory and find the index.json file
	index_file = [os.path.join(model_name, f) for f in os.listdir(model_name) if "index.json" in f]
	assert len(index_file) == 1, f"Could not find index.json in {model_name}"

	# process index json
	with open(index_file[0], "r") as f:
	index_data = json.load(f)

	bin_to_weight_mapping = dict()
	for k, v in index_data["weight_map"].items():
	if v in bin_to_weight_mapping:
	bin_to_weight_mapping[v].append(k)
	else:
	bin_to_weight_mapping[v] = [k]

	# make state dict
	state_dict = dict()
	for bin_name, weight_list in bin_to_weight_mapping.items():
	bin_path = os.path.join(model_name, bin_name)
	bin_state_dict = torch.load(bin_path, map_location=lambda storage, loc: storage)
	for weight in weight_list:
	state_dict[weight] = bin_state_dict[weight]
	return state_dict
	else:
	# if it is a file, we just load it directly in the typical PyTorch manner
	assert os.path.exists(model_name), f"Could not find DiT checkpoint at {model_name}"
	checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
	if "ema" in checkpoint: # supports checkpoints from train.py
	checkpoint = checkpoint["ema"]
	return checkpoint


	def download_model(model_name):
	"""
	Downloads a pre-trained DiT model from the web.
	"""
	assert model_name in pretrained_models
	local_path = f"pretrained_models/{model_name}"
	if not os.path.isfile(local_path):
	os.makedirs("pretrained_models", exist_ok=True)
	web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{model_name}"
	download_url(web_path, "pretrained_models")
	model = torch.load(local_path, map_location=lambda storage, loc: storage)
	return model


	if __name__ == "__main__":
	# Download all DiT checkpoints
	for model in pretrained_models:
	download_model(model)
	print("Done.")