app.py

# script to run the app in HuggingFace space
import os, argparse
import numpy as np

from HF.UI import AppUI
from HF.Utils import toPILImage
from HF.NNHelper import modelsFrom, inference_from_config
from Utils.WandBUtils import CWBProject

def infer(models):
  def _processImage(modelName, inputImage, **kwargs):
    model = models.get(modelName, None)
    assert model is not None, f'Invalid model name: {modelName}'

    assert isinstance(inputImage, np.ndarray), f'Invalid image type: {type(inputImage)}'
    assert 3 == len(inputImage.shape), f'Invalid image shape: {inputImage.shape}'
    assert np.uint8 == inputImage.dtype, f'Invalid image dtype: {inputImage.dtype}'
    assert 3 == inputImage.shape[-1], f'Invalid image channels: {inputImage.shape[-1]}'
    
    res = model(raw=inputImage, **kwargs)
    if ('video' in res):
      return dict(video=res['video'])
    
    upscaled = toPILImage(res['upscaled'], isBGR=False)
    input = toPILImage(res['input'], isBGR=False)

    return dict(upscaled=upscaled, input=input)
  return _processImage

def run2inference(run, runName=None):
  if runName is None: runName = run.name.replace('[Public] ', '')
  runConfig = run.config
  runName = "%s (%s, loss: %.5f)" % (runName, run.id, run.bestLoss)
  # add corresponding HF info
  runConfig['huggingface'] = { "name": runName, "wandb": run.fullId }
  return list(inference_from_config(runConfig))


def main(args):
  WBProject = CWBProject('green_wizard/FranNet')
  folder = os.path.dirname(os.path.abspath(__file__))
  # load list of models from the folder "configs"
  models = []
  bestGroups = WBProject.groups(onlyBest=True)
  for runName, run in bestGroups.items():
    if runName.startswith('[Public] '):
      models.extend( run2inference(run) )
    continue
  # convert to dict
  models = {model.name: model for model in models}

  app = AppUI(
    processImage=infer(models),
    models=models,
  )
  app.queue() # enable queueing of requests/events
  app.launch(inline=False, server_port=args.port, server_name=args.host)
  return

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument('--port', type=int, default=7860)
  parser.add_argument('--host', type=str, default=None)
  args = parser.parse_args()
  main(args)