feng2022's picture
Update app.py
c18ab9a
raw
history blame
1.76 kB
#!/usr/bin/env python
from __future__ import annotations
import argparse
import functools
import os
import pickle
import sys
import gradio as gr
import numpy as np
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import pipeline
sys.path.append('.')
sys.path.append('./Time_TravelRephotography')
from argparse import Namespace
from projector import (
ProjectorArguments,
main,
)
sys.path.insert(0, 'StyleGAN-Human')
TITLE = 'Time-TravelRephotography'
DESCRIPTION = '''This is an unofficial demo for https://github.com/Time-Travel-Rephotography.
'''
ARTICLE = '<center><img src="https://visitor-badge.glitch.me/badge?page_id=hysts.stylegan-human" alt="visitor badge"/></center>'
TOKEN = "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"
pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
def load_model(file_name: str, path:str,device: torch.device) -> nn.Module:
path = hf_hub_download(f'{path}',
f'{file_name}',
use_auth_token=TOKEN)
with open(path, 'rb') as f:
model = torch.load(f)
model.eval()
model.to(device)
with torch.inference_mode():
z = torch.zeros((1, model.z_dim)).to(device)
label = torch.zeros([1, model.c_dim], device=device)
model(z, label, force_fp32=True)
return model
def predict(text):
return pipe(text)[0]["translation_text"]
def main():
load_model("stylegan2-ffhq-config-f","feng2022/Time-TravelRephotography_stylegan2-ffhq-config-f")
iface = gr.Interface(
fn=predict,
inputs='text',
outputs='text',
examples=[["Time-TravelRephotography"]]
)
iface.launch()
if __name__ == '__main__':
main()