softwareweaver/MusicGen

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7 8"""9To run this script, from the root of the repo. Make sure to have Flask installed10 11    FLASK_DEBUG=1 FLASK_APP=scripts.mos flask run -p 456712    # or if you have gunicorn13    gunicorn -w 4 -b 127.0.0.1:8895 -t 120 'scripts.mos:app'  --access-logfile -14 15"""16from collections import defaultdict17from functools import wraps18from hashlib import sha119import json20import math21from pathlib import Path22import random23import typing as tp24 25from flask import Flask, redirect, render_template, request, session, url_for26 27from audiocraft import train28from audiocraft.utils.samples.manager import get_samples_for_xps29 30 31SAMPLES_PER_PAGE = 832MAX_RATING = 533storage = Path(train.main.dora.dir / 'mos_storage')34storage.mkdir(exist_ok=True)35surveys = storage / 'surveys'36surveys.mkdir(exist_ok=True)37magma_root = Path(train.__file__).parent.parent38app = Flask('mos', static_folder=str(magma_root / 'scripts/static'),39            template_folder=str(magma_root / 'scripts/templates'))40app.secret_key = b'audiocraft makes the best songs'41 42 43def normalize_path(path: Path):44    """Just to make path a bit nicer, make them relative to the Dora root dir.45    """46    path = path.resolve()47    dora_dir = train.main.dora.dir.resolve() / 'xps'48    return path.relative_to(dora_dir)49 50 51def get_full_path(normalized_path: Path):52    """Revert `normalize_path`.53    """54    return train.main.dora.dir.resolve() / 'xps' / normalized_path55 56 57def get_signature(xps: tp.List[str]):58    """Return a signature for a list of XP signatures.59    """60    return sha1(json.dumps(xps).encode()).hexdigest()[:10]61 62 63def ensure_logged(func):64    """Ensure user is logged in.65    """66    @wraps(func)67    def _wrapped(*args, **kwargs):68        user = session.get('user')69        if user is None:70            return redirect(url_for('login', redirect_to=request.url))71        return func(*args, **kwargs)72    return _wrapped73 74 75@app.route('/login', methods=['GET', 'POST'])76def login():77    """Login user if not already, then redirect.78    """79    user = session.get('user')80    if user is None:81        error = None82        if request.method == 'POST':83            user = request.form['user']84            if not user:85                error = 'User cannot be empty'86        if user is None or error:87            return render_template('login.html', error=error)88    assert user89    session['user'] = user90    redirect_to = request.args.get('redirect_to')91    if redirect_to is None:92        redirect_to = url_for('index')93    return redirect(redirect_to)94 95 96@app.route('/', methods=['GET', 'POST'])97@ensure_logged98def index():99    """Offer to create a new study.100    """101    errors = []102    if request.method == 'POST':103        xps_or_grids = [part.strip() for part in request.form['xps'].split()]104        xps = set()105        for xp_or_grid in xps_or_grids:106            xp_path = train.main.dora.dir / 'xps' / xp_or_grid107            if xp_path.exists():108                xps.add(xp_or_grid)109                continue110            grid_path = train.main.dora.dir / 'grids' / xp_or_grid111            if grid_path.exists():112                for child in grid_path.iterdir():113                    if child.is_symlink():114                        xps.add(child.name)115                continue116            errors.append(f'{xp_or_grid} is neither an XP nor a grid!')117        assert xps or errors118        blind = 'true' if request.form.get('blind') == 'on' else 'false'119        xps = list(xps)120        if not errors:121            signature = get_signature(xps)122            manifest = {123                'xps': xps,124            }125            survey_path = surveys / signature126            survey_path.mkdir(exist_ok=True)127            with open(survey_path / 'manifest.json', 'w') as f:128                json.dump(manifest, f, indent=2)129            return redirect(url_for('survey', blind=blind, signature=signature))130    return render_template('index.html', errors=errors)131 132 133@app.route('/survey/<signature>', methods=['GET', 'POST'])134@ensure_logged135def survey(signature):136    success = request.args.get('success', False)137    seed = int(request.args.get('seed', 4321))138    blind = request.args.get('blind', 'false') in ['true', 'on', 'True']139    exclude_prompted = request.args.get('exclude_prompted', 'false') in ['true', 'on', 'True']140    exclude_unprompted = request.args.get('exclude_unprompted', 'false') in ['true', 'on', 'True']141    max_epoch = int(request.args.get('max_epoch', '-1'))142    survey_path = surveys / signature143    assert survey_path.exists(), survey_path144 145    user = session['user']146    result_folder = survey_path / 'results'147    result_folder.mkdir(exist_ok=True)148    result_file = result_folder / f'{user}_{seed}.json'149 150    with open(survey_path / 'manifest.json') as f:151        manifest = json.load(f)152 153    xps = [train.main.get_xp_from_sig(xp) for xp in manifest['xps']]154    names, ref_name = train.main.get_names(xps)155 156    samples_kwargs = {157        'exclude_prompted': exclude_prompted,158        'exclude_unprompted': exclude_unprompted,159        'max_epoch': max_epoch,160    }161    matched_samples = get_samples_for_xps(xps, epoch=-1, **samples_kwargs)  # fetch latest epoch162    models_by_id = {163        id: [{164            'xp': xps[idx],165            'xp_name': names[idx],166            'model_id': f'{xps[idx].sig}-{sample.id}',167            'sample': sample,168            'is_prompted': sample.prompt is not None,169            'errors': [],170        } for idx, sample in enumerate(samples)]171        for id, samples in matched_samples.items()172    }173    experiments = [174        {'xp': xp, 'name': names[idx], 'epoch': list(matched_samples.values())[0][idx].epoch}175        for idx, xp in enumerate(xps)176    ]177 178    keys = list(matched_samples.keys())179    keys.sort()180    rng = random.Random(seed)181    rng.shuffle(keys)182    model_ids = keys[:SAMPLES_PER_PAGE]183 184    if blind:185        for key in model_ids:186            rng.shuffle(models_by_id[key])187 188    ok = True189    if request.method == 'POST':190        all_samples_results = []191        for id in model_ids:192            models = models_by_id[id]193            result = {194                'id': id,195                'is_prompted': models[0]['is_prompted'],196                'models': {}197            }198            all_samples_results.append(result)199            for model in models:200                rating = request.form[model['model_id']]201                if rating:202                    rating = int(rating)203                    assert rating <= MAX_RATING and rating >= 1204                    result['models'][model['xp'].sig] = rating205                    model['rating'] = rating206                else:207                    ok = False208                    model['errors'].append('Please rate this model.')209        if ok:210            result = {211                'results': all_samples_results,212                'seed': seed,213                'user': user,214                'blind': blind,215                'exclude_prompted': exclude_prompted,216                'exclude_unprompted': exclude_unprompted,217            }218            print(result)219            with open(result_file, 'w') as f:220                json.dump(result, f)221            seed = seed + 1222            return redirect(url_for(223                'survey', signature=signature, blind=blind, seed=seed,224                exclude_prompted=exclude_prompted, exclude_unprompted=exclude_unprompted,225                max_epoch=max_epoch, success=True))226 227    ratings = list(range(1, MAX_RATING + 1))228    return render_template(229        'survey.html', ratings=ratings, blind=blind, seed=seed, signature=signature, success=success,230        exclude_prompted=exclude_prompted, exclude_unprompted=exclude_unprompted, max_epoch=max_epoch,231        experiments=experiments, models_by_id=models_by_id, model_ids=model_ids, errors=[],232        ref_name=ref_name, already_filled=result_file.exists())233 234 235@app.route('/audio/<path:path>')236def audio(path: str):237    full_path = Path('/') / path238    assert full_path.suffix in [".mp3", ".wav"]239    return full_path.read_bytes(), {'Content-Type': 'audio/mpeg'}240 241 242def mean(x):243    return sum(x) / len(x)244 245 246def std(x):247    m = mean(x)248    return math.sqrt(sum((i - m)**2 for i in x) / len(x))249 250 251@app.route('/results/<signature>')252@ensure_logged253def results(signature):254 255    survey_path = surveys / signature256    assert survey_path.exists(), survey_path257    result_folder = survey_path / 'results'258    result_folder.mkdir(exist_ok=True)259 260    # ratings per model, then per user.261    ratings_per_model = defaultdict(list)262    users = []263    for result_file in result_folder.iterdir():264        if result_file.suffix != '.json':265            continue266        with open(result_file) as f:267            results = json.load(f)268        users.append(results['user'])269        for result in results['results']:270            for sig, rating in result['models'].items():271                ratings_per_model[sig].append(rating)272 273    fmt = '{:.2f}'274    models = []275    for model in sorted(ratings_per_model.keys()):276        ratings = ratings_per_model[model]277 278        models.append({279            'sig': model,280            'samples': len(ratings),281            'mean_rating': fmt.format(mean(ratings)),282            # the value 1.96 was probably chosen to achieve some283            # confidence interval assuming gaussianity.284            'std_rating': fmt.format(1.96 * std(ratings) / len(ratings)**0.5),285        })286    return render_template('results.html', signature=signature, models=models, users=users)287