Update app.py

app.py

@@ -15,6 +15,9 @@ import uvicorn
 import requests
 from urllib.parse import quote
 from unicodedata import normalize
+import json
+import time
+from html.parser import HTMLParser
 
 root = os.path.dirname(os.path.abspath(__file__))
 textures_folder = os.path.join(root, 'textures')
@@ -191,7 +194,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 +207,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 +246,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 +355,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,6 +370,7 @@ const PageLoadObserver = new MutationObserver((mutationsList, observer) => {
 PageLoadObserver.observe(document, { childList: true, subtree: true });
 """
 
+
 def extract_path_from_result(predict_answer):
     if isinstance(predict_answer, (tuple, list)):
         result = predict_answer[0]
@@ -449,7 +468,7 @@ async def restore_upscale(files, restore_method):
         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):
@@ -493,12 +512,168 @@ def apply_texture(input_image, textures_choice, opacity_slider):
     return [result]
 
 
+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):
+    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):
+    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)
+        status = response.json()['status']
+        if status == 'succeeded':
+            break
+        time.sleep(1)
+    session.close()
+    return response.json()['output']
+
+
+urls = [
+    'https://replicate.com/cjwbw/real-esrgan',
+    'https://replicate.com/daanelson/real-esrgan-a100',
+    'https://replicate.com/xinntao/realesrgan',
+]
+
+
+def upscaler(image):
+    img_url = upload_image(image)
+
+    def run(url):
+        try:
+            return replicate_upscale(url, img_url)
+        except Exception as e:
+            print(e)
+            return return 'https://iili.io/JzrxjDP.png'
+    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 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)
+
 with gr.Blocks(analytics_enabled=False, css=radio_css) 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=upscaler, 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"):
@@ -511,6 +686,7 @@ with gr.Blocks(analytics_enabled=False, css=radio_css) as demo:
         apply_button.click(fn=apply_texture, inputs=[input_image, textures_choice, opacity_slider], outputs=result_image, api_name="texturize")
 
 
+
 app = FastAPI()