d0tpy commited on
Commit
27a87ea
·
verified ·
1 Parent(s): ca4ec9d

Update video_enhancer.py

Browse files
Files changed (1) hide show
  1. video_enhancer.py +5 -0
video_enhancer.py CHANGED
@@ -2,6 +2,7 @@ import cv2
2
  import numpy as np
3
  import torch
4
  import io
 
5
  import asyncio
6
  from fastapi.responses import StreamingResponse
7
  from basicsr.archs.rrdbnet_arch import RRDBNet
@@ -9,6 +10,10 @@ from realesrgan import RealESRGANer
9
  from huggingface_hub import hf_hub_download
10
  from concurrent.futures import ThreadPoolExecutor
11
 
 
 
 
 
12
  class VideoEnhancer:
13
  def __init__(self, model_name="RealESRGAN_x4plus"):
14
  self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
2
  import numpy as np
3
  import torch
4
  import io
5
+ import os
6
  import asyncio
7
  from fastapi.responses import StreamingResponse
8
  from basicsr.archs.rrdbnet_arch import RRDBNet
 
10
  from huggingface_hub import hf_hub_download
11
  from concurrent.futures import ThreadPoolExecutor
12
 
13
+ cache_dir = '/tmp/huggingface_cache'
14
+ os.environ['HF_HOME'] = cache_dir
15
+ os.makedirs(cache_dir, exist_ok=True)
16
+
17
  class VideoEnhancer:
18
  def __init__(self, model_name="RealESRGAN_x4plus"):
19
  self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')