Update app.py

app.py

@@ -1,8 +1,6 @@
 import os
-os.system('pip install numpy pillow opencv-python fastapi starlette uvicorn requests')
 import gradio as gr
 import shutil
-import cv2
 import numpy as np
 from PIL import Image
 import asyncio
@@ -15,6 +13,14 @@ import uvicorn
 import requests
 from urllib.parse import quote
 from unicodedata import normalize
+import json
+import time
+from html.parser import HTMLParser
+try:
+    import cv2
+except ImportError:
+    os.system('pip install opencv-python')
+    import cv2
 
 root = os.path.dirname(os.path.abspath(__file__))
 textures_folder = os.path.join(root, 'textures')
@@ -22,7 +28,6 @@ os.makedirs(textures_folder, exist_ok=True)
 valid_extensions = ['.jpeg', '.jpg', '.png']
 
 
-
 textures_repo = "https://huggingface.co/datasets/2ch/textures/resolve/main/"
 textures_for_download = [
 f"{textures_repo}гауссовский_шум_и_мелкое_зерно.png?download=true",
@@ -50,7 +55,7 @@ f"{textures_repo}шумная_матрица.png?download=true",
     ]
 
 
-def dl_textures(texture_url):
+def dl_textures(texture_url: str) -> None:
     texture_for_download = quote(normalize('NFD', texture_url), safe='/?:=')
     filename = texture_url.split('/')[-1].split('?')[0]
     file_path = os.path.join(textures_folder, filename)
@@ -61,7 +66,7 @@ def dl_textures(texture_url):
             f.write(chunk)
 
 
-def create_texture_preview(texture_folder, output_folder, size=(246, 246)):
+def create_texture_preview(texture_folder: str, output_folder: str, size: tuple[int] = (246, 246)) -> None:
     os.makedirs(output_folder, exist_ok=True)
     for texture in os.listdir(texture_folder):
         img_path = os.path.join(texture_folder, texture)
@@ -72,21 +77,17 @@ def create_texture_preview(texture_folder, output_folder, size=(246, 246)):
         cv2.imwrite(os.path.join(output_folder, texture), img)
 
 
-def prepare_textures(texture_folder, output_folder):
+def prepare_textures(texture_folder: str, output_folder: str) -> None:
     with ThreadPoolExecutor(max_workers=len(textures_for_download)) as executor:
         futures = [executor.submit(dl_textures, texture_for_download) for texture_for_download in
                    textures_for_download]
         for future in as_completed(futures):
             future.result()
-
     create_texture_preview(texture_folder, output_folder, size=(246, 246))
 
 
-prepare_textures(textures_folder, os.path.join(root, 'preview'))
-
-
-
 preview_css = ""
+prepare_textures(textures_folder, os.path.join(root, 'preview'))
 for i, texture in enumerate(os.listdir(textures_folder), start=1):
     if os.path.splitext(texture)[1].lower() in valid_extensions:
         preview_css += f"""[data-testid="{i:02d}-radio-label"]::before {{
@@ -191,7 +192,8 @@ footer:after {
 }
 
 #textured_result-download-link,
-#restored_image-download-link {
+#restored_image-download-link,
+#upscaled_image-download-link {
     position: absolute;
     z-index: 9999;
     padding: 2px 4px;
@@ -203,17 +205,21 @@ footer:after {
     transition: 300ms
 }
 
-#download-link:hover {
+#textured_result-download-link:hover,
+#restored_image-download-link:hover,
+#upscaled_image-download-link:hover {
     color: #99f7a8
 }
 
-#restored_images.disabled {
+#restored_images.disabled,
+#upscaled_images.disabled {
     height: 0px !important;
     opacity: 0;
     transition: 300ms
 }
 
-#restored_images.enabled {
+#restored_images.enabled,
+#upscaled_images.enabled {
     transition: 300ms
 }
 """ + preview_css
@@ -238,24 +244,28 @@ const PageLoadObserver = new MutationObserver((mutationsList, observer) => {
                         document.querySelector("label[data-testid='05-radio-label']").click()
                     }, 150);
                 })
-                let RestoredGallery = document.getElementById('restored_images');
-                function checkImagesAndSetClass() {
-                    const firstDiv = RestoredGallery.querySelector('div:first-child');
+                function checkImagesAndSetClass(galleryElement) {
+                    const firstDiv = galleryElement.querySelector('div:first-child');
                     const hasChildElements = firstDiv && firstDiv.children.length > 0;
-                    const hasImages = RestoredGallery.querySelectorAll('img').length > 0;
+                    const hasImages = galleryElement.querySelectorAll('img').length > 0;
                     if (hasChildElements || hasImages) {
-                        RestoredGallery.classList.add('enabled');
-                        RestoredGallery.classList.remove('disabled');
+                        galleryElement.classList.add('enabled');
+                        galleryElement.classList.remove('disabled');
                     } else {
-                        RestoredGallery.classList.add('disabled');
-                        RestoredGallery.classList.remove('enabled');
+                        galleryElement.classList.add('disabled');
+                        galleryElement.classList.remove('enabled');
                     }
                 }
-                const FaceResoreResultCheck = new MutationObserver((mutations) => {
-                checkImagesAndSetClass();
-                });
-                FaceResoreResultCheck.observe(RestoredGallery, {childList: true, subtree: true});
-                checkImagesAndSetClass();
+                function setupGalleryObserver(galleryId) {
+                    let gallery = document.getElementById(galleryId);
+                    const observer = new MutationObserver(() => {
+                        checkImagesAndSetClass(gallery);
+                    });
+                    observer.observe(gallery, { childList: true, subtree: true });
+                    checkImagesAndSetClass(gallery);
+                }
+                setupGalleryObserver('restored_images');
+                setupGalleryObserver('upscaled_images');
                 function magnify(imgID, zoom) {
                     var img, glass, w, h, bw;
                     img = document.querySelector(imgID);
@@ -343,6 +353,12 @@ const PageLoadObserver = new MutationObserver((mutationsList, observer) => {
                   });
 
                   ImageRestoredObserver.observe(document, DownloadLinkObserverOptions);
+                  
+                  const ImageUpscaledObserver = new MutationObserver((mutationsList, observer) => {
+                    DownloadLinkObserverCallback(mutationsList, observer, '#upscaled_images img[data-testid="detailed-image"]', '#upscaled_image-download-link', 'upscaled_image-download-link');
+                  });
+
+                  ImageUpscaledObserver.observe(document, DownloadLinkObserverOptions);
 
             }
         }
@@ -352,7 +368,8 @@ const PageLoadObserver = new MutationObserver((mutationsList, observer) => {
 PageLoadObserver.observe(document, { childList: true, subtree: true });
 """
 
-def extract_path_from_result(predict_answer):
+
+def extract_path_from_result(predict_answer: str | list[str] | tuple[str]) -> str:
     if isinstance(predict_answer, (tuple, list)):
         result = predict_answer[0]
         shutil.rmtree(os.path.dirname(predict_answer[1]), ignore_errors=True)
@@ -386,7 +403,6 @@ def restore_face_codeformer(img_path: str) -> None:
     restore_face_common(img_path, result, "codeformer")
 
 
-
 async def restore_faces_one_image(img_path: str, func_list: list) -> bool:
     def run_func(func) -> bool:
         for _ in range(3):
@@ -435,7 +451,7 @@ def get_file_paths(input_path: str | list[str], extensions_list: list[str]) -> l
     return files
 
 
-async def restore_upscale(files, restore_method):
+async def restore_upscale(files: tuple, restore_method: str) -> list[str]:
     file_paths = [file.name for file in files]
     if restore_method == 'codeformer':
         func_list = [restore_face_codeformer]
@@ -446,10 +462,9 @@ async def restore_upscale(files, restore_method):
     results = await restore_faces_batch(file_paths, func_list, batch_size=3)
     if results:
         file_paths = get_file_paths(file_paths, valid_extensions)
-        print(f"restore_upscale: get_file_paths: {file_paths}")
         return file_paths
     else:
-        return [os.path.join(root, 'error.png')]
+        return ['https://iili.io/JzrxjDP.png']
 
 
 def image_noise_softlight_layer_mix(img, texture, output: str = None, opacity: float = 0.7):
@@ -488,17 +503,185 @@ def image_noise_softlight_layer_mix(img, texture, output: str = None, opacity: f
     return np.array(image)
 
 
-def apply_texture(input_image, textures_choice, opacity_slider):
+def apply_texture(input_image, textures_choice: str, opacity_slider: float):
     result = image_noise_softlight_layer_mix(input_image, os.path.join(textures_folder, textures_choice), opacity=opacity_slider)
     return [result]
 
 
-with gr.Blocks(analytics_enabled=False, css=radio_css) as demo:
+def temp_upload_file(file_path: str) -> str | None:
+    servers = [
+        ('https://transfer.sh/', 'fileToUpload'),
+        ('https://x0.at/', 'file'),
+        ('https://tmpfiles.org/api/v1/upload', 'file'),
+        ('https://uguu.se/upload.php', 'files[]')
+    ]
+    for i in range(3):
+        for server, file_key in servers:
+            try:
+                with open(file_path, 'rb') as f:
+                    files = {file_key: f}
+                    response = requests.post(server, files=files)
+                    if response.status_code == 200:
+                        if server == 'https://transfer.sh/':
+                            return response.text.replace("https://transfer.sh/","https://transfer.sh/get/").replace("\n","")
+                        elif server == 'https://tmpfiles.org/api/v1/upload':
+                            response_json = response.json()
+                            if response_json['status'] == 'success':
+                                return response_json['data']['url'].replace("https://tmpfiles.org/", "https://tmpfiles.org/dl/")
+                        elif server == 'https://uguu.se/upload.php':
+                            response_json = response.json()
+                            if response_json['success']:
+                                return response_json['files'][0]['url']
+                        else:
+                            return response.text
+            except Exception as e:
+                print(f'{server}: {e}')
+    return None
+
+
+def upload_image(image_path: str) -> str | None:
+    files = {'source': open(image_path, "rb")}
+    data = {'key': '6d207e02198a847aa98d0a2a901485a5', 'action': 'upload', 'format': 'json'}
+    response = requests.post('https://freeimage.host/api/1/upload', files=files, data=data)
+    if response.json()["status_code"] == 200:
+        return response.json()["image"]["url"]
+    else:
+        return temp_upload_file(image_path)
+
+
+def get_headers(url: str) -> dict:
+    session = requests.Session()
+    anon_auth = session.get(url)
+    cookies = session.cookies.get_dict()
+    return {
+        'content-type': 'application/json',
+        'cookie': f'csrftoken={cookies["csrftoken"]}; replicate_anonymous_id={cookies["replicate_anonymous_id"]};',
+        'origin': 'https://replicate.com',
+        'x-csrftoken': cookies['csrftoken'],
+        'authority': 'replicate.com',
+        'accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/jxl,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
+        'accept-language': 'ru-RU,ru;q=0.9,en-US;q=0.8,en;q=0.7',
+        'cache-control': 'no-cache',
+        'dnt': '1',
+        'pragma': 'no-cache',
+        'referer': f'{url}?input=http',
+        'sec-ch-ua': '"Chromium";v="117", "Not;A=Brand";v="8"',
+        'sec-ch-ua-mobile': '?0',
+        'sec-ch-ua-platform': '"Windows"',
+        'sec-fetch-dest': 'document',
+        'sec-fetch-mode': 'navigate',
+        'sec-fetch-site': 'same-origin',
+        'sec-fetch-user': '?1',
+        'upgrade-insecure-requests': '1',
+        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/117.0.0.0 Safari/537.36'
+    }
+
+
+def get_version(url: str) -> str:
+    url = url.rstrip('/') + '/versions'
+    response = requests.get(url)
+
+    class Version(HTMLParser):
+        def __init__(self):
+            super().__init__()
+            self.recording = 0
+            self.data = ''
+
+        def handle_starttag(self, tag, attrs):
+            if tag == 'a':
+                for name, value in attrs:
+                    if name == 'href' and '/versions/' in value:
+                        self.recording = 1
+
+        def handle_endtag(self, tag):
+            if tag == 'a' and self.recording:
+                self.recording -= 1
+
+        def handle_data(self, data):
+            if self.recording:
+                self.data = data
+
+    parser = Version()
+    parser.feed(response.text)
+    return parser.data.strip()
+
+
+def replicate_upscale(url: str, image_url: str, upscale: int = 2) -> str:
+    version = get_version(url)
+    headers = get_headers(url)
+    session = requests.Session()
+    anon_auth = session.get(url, headers=headers)
+    data = {
+        "version": version,
+        "input": {
+            "img": image_url,
+            "image": image_url,
+            "upscale": upscale,
+            "scale": upscale,
+            "version": "General - RealESRGANplus",
+        },
+        "face_enhance": False,
+        "is_training": False,
+        "stream": False
+    }
+    response = session.post('https://replicate.com/api/predictions', headers=headers, data=json.dumps(data))
+    prediction_id = response.json()['id']
+    while True:
+        response = session.get(f'https://replicate.com/api/predictions/{prediction_id}', headers=headers)
+        if 'status' in response.json():
+            status = response.json()['status']
+        else:
+            status = 'processing'
+        if status == 'succeeded':
+            break
+        time.sleep(1)
+    session.close()
+    return response.json()['output']
+
+
+def upscaler(img_url: str) -> list[str] | None:
+    def run(url):
+        try:
+            return replicate_upscale(url, img_url)
+        except Exception as e:
+            print(e)
+            return None
+
+    urls = [
+        'https://replicate.com/cjwbw/real-esrgan',
+        'https://replicate.com/daanelson/real-esrgan-a100',
+        'https://replicate.com/xinntao/realesrgan',
+    ]
+    with ThreadPoolExecutor() as executor:
+        futures = {executor.submit(run, url) for url in urls}
+        for future in as_completed(futures):
+            result = future.result()
+            if result is not None:
+                break
+    return [result]
+
+
+def check_upscale_result(image: str) -> list[str]:
+    attempt = 0
+    response = None
+    while attempt < 3:
+        response = upscaler(upload_image(image))
+        if response:
+            return response
+        attempt += 1
+    return ['https://iili.io/JzrxjDP.png']
+
+
+with gr.Blocks(analytics_enabled=False, css=radio_css, theme='Taithrah/Minimal', title='апскейл') as demo:
+    with gr.Tab(label="апскейл", elem_id="upscale_tab"):
+        file_output = gr.Gallery(label="", container=True, object_fit="cover", columns=4, rows=4, allow_preview=True, preview=True, show_share_button=False, show_download_button=False, elem_id="upscaled_images")
+        upload_button = gr.UploadButton("выбор одного изображения для обработки", file_types=["image"], file_count="single", variant="primary")
+        upload_button.upload(fn=check_upscale_result, inputs=[upload_button], outputs=file_output, api_name="upscale")
     with gr.Tab(label="восстановление лиц", id=1, elem_id="restore_tab"):
         restore_method = gr.Radio(["codeformer", "gfpgan", "оба"], value="codeformer", label="", interactive=True)
         restore_method.change(fn=lambda x: print(f"restore_method value = {x}"), inputs=restore_method, api_name="show_selected_method")
         file_output = gr.Gallery(label="", container=True, object_fit="cover", columns=4, rows=4, allow_preview=True, preview=True, show_share_button=False, show_download_button=False, elem_id="restored_images")
-        upload_button = gr.UploadButton("выбор изображений для обработки", file_types=["image"], file_count="multiple", variant="primary")
+        upload_button = gr.UploadButton("выбор нескольких изображений для обработки", file_types=["image"], file_count="multiple", variant="primary")
         upload_button.upload(fn=restore_upscale, inputs=[upload_button, restore_method], outputs=file_output, api_name="face_restore")
     with gr.Tab(label="наложение зернистости пленки и шума", id=2, elem_id="textures_tab"):
         with gr.Row(variant="compact", elem_id="textures_tab_images"):