Initial commit

osu_dreamer_gui/gui/handlers/mp3_choose_upload.py

@@ -0,0 +1,29 @@
+from nicegui import ui, events, app
+import mutagen
+
+from ...modules.encoder.encoder import dump
+
+
+async def on_upload(e: events.UploadEventArguments):
+    storage = app.storage.user
+
+    storage['bpm'] = None
+    storage['detected_title'] = None
+    storage['detected_artist'] = None
+
+    tags = mutagen.File(e.content, easy=True)
+
+    if 'bpm' in tags:
+        storage['bpm'] = int(float(tags['bpm'][0]))
+
+    if 'title' in tags:
+        storage['detected_title'] = tags['title'][0]
+
+    if 'artist' in tags:
+        storage['detected_artist'] = tags['artist'][0]
+
+    storage['filename'] = e.name
+    storage['audio_content'] = dump(e.content.read())
+
+    storage['can_be_created'] = True
+    storage['can_be_saved'] = False

osu_dreamer_gui/gui/handlers/on_create.py

@@ -0,0 +1,57 @@
+from nicegui import ui, app, run
+from io import BytesIO
+from ...modules.encoder.encoder import load
+
+
+def creator(storage):
+    import torch
+    from osu_dreamer.model import Model
+    from ...modules.generate import generate_mapset
+    from rich.traceback import install
+
+    install(show_locals=True)
+
+    model = Model.load_from_checkpoint(
+        f'models/{storage["model_name"]}',
+        sample_steps=storage['sample_steps'],
+    ).eval()
+
+    if torch.cuda.is_available():
+        model = model.cuda()
+
+    audio = BytesIO(load(storage['audio_content']))
+    audio.name = storage['filename']
+    audio.seek(0)
+    return generate_mapset(
+        model,
+        audio,
+        storage['bpm'],
+        storage['num_samples'],
+        storage['detected_title'], storage['detected_artist'],
+    ).name
+
+
+async def on_create():
+    storage = app.storage.user
+
+    storage['is_loading'] = True
+    storage['can_be_created'] = False
+
+    try:
+        mapset_path = await run.cpu_bound(creator, dict(
+            model_name=storage['model_name'],
+            audio_content=storage['audio_content'],
+            bpm=storage['bpm'],
+            num_samples=storage['num_samples'],
+            sample_steps=storage['sample_steps'],
+            detected_title=storage['detected_title'],
+            detected_artist=storage['detected_artist'],
+            filename=storage['filename'],
+        ))
+        storage['mapset_path'] = mapset_path
+        storage['can_be_saved'] = True
+    except Exception as e:
+        ui.notify(f'Error {e}')
+
+    storage['is_loading'] = False
+    storage['can_be_created'] = True