Update app.py

app.py

@@ -7,6 +7,34 @@ import os
 from io import BytesIO
 import html
 import re
+import cv2
+import torch
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from gfpgan.utils import GFPGANer
+from realesrgan.utils import RealESRGANer
+
+os.system("pip freeze")
+# download weights
+if not os.path.exists('realesr-general-x4v3.pth'):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+if not os.path.exists('GFPGANv1.2.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth -P .")
+if not os.path.exists('GFPGANv1.3.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .")
+if not os.path.exists('GFPGANv1.4.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('RestoreFormer.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
+if not os.path.exists('CodeFormer.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth -P .")
+
+# background enhancer with RealESRGAN
+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)
+
+os.makedirs('output', exist_ok=True)
 
 
 
@@ -174,25 +202,59 @@ def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, he
     return job["imageUrl"]
 
 
-def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed):
-    result = prodia_client.transform({
-        "imageData": image_to_base64(input_image),
-        "denoising_strength": denoising,
-        "prompt": prompt,
-        "negative_prompt": negative_prompt,
-        "model": model,
-        "steps": steps,
-        "sampler": sampler,
-        "cfg_scale": cfg_scale,
-        "width": width,
-        "height": height,
-        "seed": seed
-    })
-
-    job = prodia_client.wait(result)
-
-    return job["imageUrl"]
+def img2img(img, version, scale, weight):
+    weight /= 100
+    print(img, version, scale, weight)
+    try:
+        extension = os.path.splitext(os.path.basename(str(img)))[1]
+        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            img_mode = 'RGBA'
+        elif len(img.shape) == 2:  # for gray inputs
+            img_mode = None
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        else:
+            img_mode = None
+
+        if version == 'v1.2':
+            face_enhancer = GFPGANer(
+            model_path='GFPGANv1.2.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'v1.3':
+            face_enhancer = GFPGANer(
+            model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'v1.4':
+            face_enhancer = GFPGANer(
+            model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'RestoreFormer':
+            face_enhancer = GFPGANer(
+            model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'CodeFormer':
+            face_enhancer = GFPGANer(
+            model_path='CodeFormer.pth', upscale=2, arch='CodeFormer', channel_multiplier=2, bg_upsampler=upsampler)
+
+        try:
+            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
+        except RuntimeError as error:
+            print('Error', error)
+
+        try:
+            interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+            h, w = img.shape[0:2]
+            output = cv2.resize(output, (int(w * scale), int(h * scale)), interpolation=interpolation)
+        except Exception as error:
+            print('wrong scale input.', error)
+        if img_mode == 'RGBA':  # RGBA images should be saved in png format
+            extension = 'png'
+        else:
+            extension = 'jpg'
+        save_path = f'output/out.{extension}'
+        cv2.imwrite(save_path, output)
 
+        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        return output
+    except Exception as error:
+        print('global exception', error)
+        return None
 
 css = """
 footer {visibility: hidden !important;}
@@ -224,11 +286,24 @@ with gr.Blocks(css=css) as demo:
                     
                         cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7, step=1)
                         seed = gr.Slider(label="Seed", minimum=-1, maximum=10000000, value=-1)
+                with gr.Tab("Настройки апскейлинга"):
+                    with gr.Row():
+                        version = gr.inputs.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer', 'CodeFormer'], type="value", default='v1.4', label='Версия'),
+                        scale = gr.inputs.Number(label="Коэффициент масштабирования", default=2),
+                        weight = gr.Slider(0, 100, label='Weight, только для CodeFormer. 0 для лучшего качества, 100 для лучшей идентичности', default=50)
+
 
             with gr.Column():
                 text_button = gr.Button("Создать", variant='primary', elem_id="generate")
             with gr.Column(scale=2):
-                image_output = gr.Image()
+                image_output = gr.Image(show_label=True, label='Сгенерированное изображение:')
+
+            with gr.Column():
+                text_button_up = gr.Button("Улучшить качество", variant='secondary', elem_id="upscb")
+            with gr.Column(scale=2):
+                image_output_up = gr.Image(show_label=True, label='Увеличенное изображение:')
+
         text_button.click(txt2img, inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed], outputs=image_output)
+        text_button_up.click(img2img, inputs=[image_output, version, scale, weight], outputs=image_output_up)
         
 demo.queue(concurrency_count=64, max_size=80, api_open=False).launch(max_threads=256)