test gfpgan

app_onediff.py

@@ -4,6 +4,7 @@ import time
 import torch
 import os
 import json
+import subprocess
 
 from diffusers import (
     DDPMScheduler,
@@ -11,6 +12,20 @@ from diffusers import (
     AutoencoderKL,
 )
 
+def runcmd(cmd, verbose = False, *args, **kwargs):
+
+    process = subprocess.Popen(
+        cmd,
+        stdout = subprocess.PIPE,
+        stderr = subprocess.PIPE,
+        text = True,
+        shell = True
+    )
+    std_out, std_err = process.communicate()
+    if verbose:
+        print(std_out.strip(), std_err)
+    pass
+
 os.system("python3 -m pip --no-cache-dir install --pre nexfort -f https://github.com/siliconflow/nexfort_releases/releases/expanded_assets/torch2.4.1_cu121")
 os.system("git clone https://github.com/siliconflow/onediff.git")
 os.system("cd onediff && python3 -m pip install .")
@@ -44,7 +59,23 @@ base_pipe = AutoPipelineForText2Image.from_pretrained(
 )
 base_pipe.to(device)
 
-base_pipe = nexfort_compile(base_pipe)
+# base_pipe = nexfort_compile(base_pipe)
+
+from gfpgan.utils import GFPGANer
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from realesrgan.utils import RealESRGANer
+
+if not os.path.exists('GFPGANv1.4.pth'):
+    runcmd("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    runcmd("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
 
 def create_demo() -> gr.Blocks:
 

requirements.txt

@@ -5,3 +5,7 @@ diffusers
 transformers
 accelerate
 spaces
+gfpgan
+facexlib
+realesrgan
+basicsr