|
from pathlib import Path
|
|
|
|
import hydra
|
|
import torch
|
|
from omegaconf import DictConfig
|
|
from slider import Beatmap
|
|
|
|
from osudiffusion import DiT_models
|
|
from osuT5.inference import Preprocessor, Pipeline, Postprocessor, DiffisionPipeline
|
|
from osuT5.tokenizer import Tokenizer
|
|
from osuT5.utils import get_model
|
|
|
|
|
|
def get_args_from_beatmap(args: DictConfig):
|
|
if args.beatmap_path is None or args.beatmap_path == "":
|
|
return
|
|
|
|
beatmap_path = Path(args.beatmap_path)
|
|
|
|
if not beatmap_path.is_file():
|
|
raise FileNotFoundError(f"Beatmap file {beatmap_path} not found.")
|
|
|
|
beatmap = Beatmap.from_path(beatmap_path)
|
|
args.audio_path = beatmap_path.parent / beatmap.audio_filename
|
|
args.output_path = beatmap_path.parent
|
|
args.bpm = beatmap.bpm_max()
|
|
args.offset = min(tp.offset.total_seconds() * 1000 for tp in beatmap.timing_points)
|
|
args.slider_multiplier = beatmap.slider_multiplier
|
|
args.title = beatmap.title
|
|
args.artist = beatmap.artist
|
|
args.beatmap_id = beatmap.beatmap_id if args.beatmap_id == -1 else args.beatmap_id
|
|
args.diffusion.style_id = beatmap.beatmap_id if args.diffusion.style_id == -1 else args.diffusion.style_id
|
|
args.difficulty = float(beatmap.stars()) if args.difficulty == -1 else args.difficulty
|
|
|
|
|
|
def find_model(ckpt_path, args: DictConfig, device):
|
|
assert Path(ckpt_path).exists(), f"Could not find DiT checkpoint at {ckpt_path}"
|
|
checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
|
|
if "ema" in checkpoint:
|
|
checkpoint = checkpoint["ema"]
|
|
|
|
model = DiT_models[args.diffusion.model](
|
|
num_classes=args.diffusion.num_classes,
|
|
context_size=19 - 3 + 128,
|
|
).to(device)
|
|
model.load_state_dict(checkpoint)
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
@hydra.main(config_path="configs", config_name="inference", version_base="1.1")
|
|
def main(args: DictConfig):
|
|
get_args_from_beatmap(args)
|
|
|
|
torch.set_grad_enabled(False)
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
ckpt_path = Path(args.model_path)
|
|
model_state = torch.load(ckpt_path / "pytorch_model.bin", map_location=device)
|
|
tokenizer_state = torch.load(ckpt_path / "custom_checkpoint_0.pkl")
|
|
|
|
tokenizer = Tokenizer()
|
|
tokenizer.load_state_dict(tokenizer_state)
|
|
|
|
model = get_model(args, tokenizer)
|
|
model.load_state_dict(model_state)
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
preprocessor = Preprocessor(args)
|
|
audio = preprocessor.load(args.audio_path)
|
|
sequences = preprocessor.segment(audio)
|
|
total_duration_ms = len(audio) / 16000 * 1000
|
|
args.total_duration_ms = total_duration_ms
|
|
|
|
|
|
|
|
|
|
|
|
generated_maps = []
|
|
generated_positions = []
|
|
diffs = []
|
|
|
|
|
|
if args.full_set:
|
|
for i in range(args.set_difficulties):
|
|
diffs.append(3 + i * (7 - 3) / (args.set_difficulties - 1))
|
|
|
|
print(diffs)
|
|
for diff in diffs:
|
|
print(f"Generating difficulty {diff}")
|
|
args.difficulty = diff
|
|
pipeline = Pipeline(args, tokenizer)
|
|
events = pipeline.generate(model, sequences)
|
|
generated_maps.append(events)
|
|
else:
|
|
pipeline = Pipeline(args, tokenizer)
|
|
events = pipeline.generate(model, sequences)
|
|
generated_maps.append(events)
|
|
|
|
|
|
|
|
if args.generate_positions:
|
|
model = find_model(args.diff_ckpt, args, device)
|
|
refine_model = find_model(args.diff_refine_ckpt, args, device) if len(args.diff_refine_ckpt) > 0 else None
|
|
diffusion_pipeline = DiffisionPipeline(args.diffusion)
|
|
for events in generated_maps:
|
|
events = diffusion_pipeline.generate(model, events, refine_model)
|
|
generated_positions.append(events)
|
|
else:
|
|
generated_positions = generated_maps
|
|
|
|
postprocessor = Postprocessor(args)
|
|
postprocessor.generate(generated_positions)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
