Spaces:
Running
Running
| import logging | |
| import os | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_url, hf_hub_download | |
| from inference.face_detector import StatRetinaFaceDetector | |
| from inference.model_pipeline import VSNetModelPipeline | |
| from inference.onnx_model import ONNXModel | |
| logging.basicConfig( | |
| format='%(asctime)s %(levelname)-8s %(message)s', | |
| level=logging.INFO, | |
| datefmt='%Y-%m-%d %H:%M:%S') | |
| MODEL_IMG_SIZE = 512 | |
| usage_count = 82 # Based on hugging face logs | |
| def load_model(): | |
| REPO_ID = "Podtekatel/ArcaneVSK2" | |
| FILENAME_OLD = "arcane_exp_230_ep_136_512_res_V2_lighter.onnx" | |
| global model_old | |
| global pipeline_old | |
| # Old model | |
| model_path = hf_hub_download(REPO_ID, FILENAME_OLD, use_auth_token=os.getenv('HF_TOKEN')) | |
| model_old = ONNXModel(model_path) | |
| pipeline_old = VSNetModelPipeline(model_old, StatRetinaFaceDetector(MODEL_IMG_SIZE), background_resize=1024, no_detected_resize=1024) | |
| return model_old | |
| load_model() | |
| def inference(img): | |
| img = np.array(img) | |
| out_img = pipeline_old(img) | |
| out_img = Image.fromarray(out_img) | |
| global usage_count | |
| usage_count += 1 | |
| logging.info(f'Usage count is {usage_count}') | |
| return out_img | |
| title = "ARCNStyleTransferV2" | |
| description = "Gradio Demo for Arcane Season 1 style transfer. To use it, simply upload your image, or click one of the examples to load them. Press ❤️ if you like this space!" | |
| article = "This is one of my successful experiments on style transfer. I've built my own pipeline, generator model and private dataset to train this model<br>" \ | |
| "" \ | |
| "" \ | |
| "" \ | |
| "Model pipeline which used in project is improved CartoonGAN.<br>" \ | |
| "This model was trained on RTX 2080 Ti 3 days with batch size 7.<br>" \ | |
| "Model weights 80 MB in ONNX fp32 format, infers 100 ms on GPU and 600 ms on CPU at 512x512 resolution.<br>" \ | |
| "If you want to use this app or integrate this model into yours, please contact me at email 'neuromancer.ai.lover@gmail.com'." | |
| imgs_folder = 'demo' | |
| examples = [[os.path.join(imgs_folder, img_filename)] for img_filename in sorted(os.listdir(imgs_folder))] | |
| demo = gr.Interface( | |
| fn=inference, | |
| inputs=[gr.inputs.Image(type="pil")], | |
| outputs=gr.outputs.Image(type="pil"), | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples) | |
| demo.queue(concurrency_count=1) | |
| demo.launch() | |