Integrate servertask into demo app, add dockerfile

.gitignore

@@ -0,0 +1,2 @@
+worker/img_received
+worker/result

Dockerfile_worker

@@ -0,0 +1,14 @@
+FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-runtime
+
+WORKDIR /usr/app
+ADD worker/requirements.txt .
+RUN pip install -r requirements.txt
+
+ADD wise .
+
+WORKDIR /usr/app/worker
+ADD  worker/serve.py .
+
+EXPOSE 8600
+
+CMD ["python", "serve.py"] 

Whitebox_style_transfer.py

@@ -9,6 +9,7 @@ import requests
 import torch
 import torch.nn.functional as F
 from PIL import Image
+import time
 
 PACKAGE_PARENT = 'wise'
 SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
@@ -144,34 +145,15 @@ def optimize(effect, preset, result_image_placeholder):
     content = st.session_state["Content_im"]
     style = st.session_state["Style_im"]
     result_image_placeholder.text("<- Custom content/style needs to be style transferred")
-    st.sidebar.info("Note: Optimizing takes up to 5 minutes.")
+    st.sidebar.warning("Note: Optimizing takes up to 5 minutes.")
     optimize_button = st.sidebar.button("Optimize Style Transfer")
     if optimize_button:
-        if HUGGING_FACE:
-            result_image_placeholder.warning("NST optimization is currently disabled in this HuggingFace Space because it takes ~5min to optimize. To try it out, please clone the repo and change the huggingface variable in demo_config.py")
-            st.stop()
-
-        result_image_placeholder.text("Executing NST to create reference image..")
-        base_dir = f"result/{datetime.datetime.now().strftime(r'%Y-%m-%d %H.%Mh %Ss')}"
-        os.makedirs(base_dir)
-        with st.spinner(text="Running NST"):
-            reference = strotss(pil_resize_long_edge_to(content, 1024),
-                                pil_resize_long_edge_to(style, 1024), content_weight=16.0,
-                                device=torch.device("cuda"), space="uniform")
-        progress_bar = result_image_placeholder.progress(0.0)
-        ref_save_path = os.path.join(base_dir, "reference.jpg")
-        content_save_path = os.path.join(base_dir, "content.jpg")
-        resize_to = 720
-        reference = pil_resize_long_edge_to(reference, resize_to)
-        reference.save(ref_save_path)
-        content.save(content_save_path)
-        ST_CONFIG["n_iterations"] = 300
         with st.spinner(text="Optimizing parameters.."):
-            vp, content_img_cuda = single_optimize(effect, preset, "l1", content_save_path, str(ref_save_path),
-                                              write_video=False, base_dir=base_dir,
-                                              iter_callback=lambda i: progress_bar.progress(
-                                                  float(i) / ST_CONFIG["n_iterations"]))
-        return content_img_cuda.detach(), vp.cuda().detach()
+            if HUGGING_FACE:
+                optimize_on_server(content, style, result_image_placeholder)
+            else:
+                optimize_params(effect, preset, content, style, result_image_placeholder)
+        return st.session_state["effect_input"], st.session_state["result_vp"]
     else:
         if not "result_vp" in st.session_state:
             st.stop()
@@ -224,6 +206,15 @@ coll2.header("Global Edits")
 result_image_placeholder = coll1.empty()
 result_image_placeholder.markdown("## loading..")
 
+from tasks import optimize_on_server, optimize_params, monitor_task
+
+if "current_server_task_id" not in st.session_state:
+    st.session_state['current_server_task_id'] = None
+
+if HUGGING_FACE and st.session_state['current_server_task_id'] is not None: 
+    with st.spinner(text="Optimizing parameters.."):
+        monitor_task(result_image_placeholder)
+
 img_choice_panel("Content", content_urls, "portrait", expanded=True)
 img_choice_panel("Style", style_urls, "starry_night", expanded=True)
 

demo_config.py

@@ -1 +1,2 @@
-HUGGING_FACE=True # if run in hugging face. Disables some things like full NST optimization
+HUGGING_FACE=True # if run in hugging face. Huggingface uses extra server task for optim
+WORKER_URL="http://mr2632.byod.hpi.de:8600"

pages/test.py

@@ -1,119 +0,0 @@
-import base64
-import datetime
-import os
-import sys
-from io import BytesIO
-from pathlib import Path
-import numpy as np
-import requests
-import torch
-import torch.nn.functional as F
-from PIL import Image
-import json
-import time
-
-PACKAGE_PARENT = 'wise'
-SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
-sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
-
-import streamlit as st
-from streamlit.logger import get_logger
-from st_click_detector import click_detector
-import streamlit.components.v1 as components
-from streamlit.source_util import get_pages
-from streamlit_extras.switch_page_button import switch_page
-
-import helpers.session_state as session_state 
-from parameter_optimization.strotss_org import strotss, pil_resize_long_edge_to
-from helpers import torch_to_np, np_to_torch
-from effects import get_default_settings, MinimalPipelineEffect 
-
-st.set_page_config(layout="wide")
-BASE_URL = "https://ivpg.hpi3d.de/wise/wise-demo/images/"
-LOGGER = get_logger(__name__)
-
-
-def upload_form(imgtype):
-    with st.form(imgtype + "-form", clear_on_submit=True):
-        uploaded_im = st.file_uploader(f"Load {imgtype} image:", type=["png", "jpg"], )
-        upload_pressed = st.form_submit_button("Upload")
-
-        if upload_pressed and uploaded_im is not None:
-            img = Image.open(uploaded_im)
-            buffered = BytesIO()
-            img.save(buffered, format="JPEG")
-            encoded = base64.b64encode(buffered.getvalue()).decode()
-            # session_state.get(uploaded_im=img, content_render_src=f"data:image/jpeg;base64,{encoded}")
-            session_state.get(**{f"{imgtype}_im": img, f"{imgtype}_render_src": f"data:image/jpeg;base64,{encoded}",
-                                    f"{imgtype}_id": "uploaded"})
-            st.session_state["action"] = "uploaded"
-            st.write("uploaded.")
-
-upload_form("Content")
-upload_form("Style")
-content = st.session_state["Content_im"]
-style = st.session_state["Style_im"]
-base_url = "http://mr2632.byod.hpi.de:5000"
-
-if content is not None and style is not None:
-    optimize_button = st.sidebar.button("Optimize Style Transfer")
-    if optimize_button:
-        url = base_url + "/upload"
-        content_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
-        style_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
-        content = pil_resize_long_edge_to(content, 1024)
-        content.save(content_path)
-        style = pil_resize_long_edge_to(style, 1024)
-        style.save(style_path)
-        files = {'style-image': open(style_path, "rb"), "content-image": open(content_path, "rb")}
-        print("start-optimizing")
-        task_id_res = requests.post(url, files=files)
-        if task_id_res.status_code != 200:
-            st.error(task_id_res.content)
-            st.stop()
-        else:
-            task_id = task_id_res.json()['task_id']
-        
-        progress_bar = st.empty()
-        with st.spinner(text="Optimizing parameters.."):
-            started_time = time.time()
-            while True:
-                time.sleep(3)
-                status = requests.get(base_url+"/get_status", params={"task_id": task_id})
-                if status.status_code != 200:
-                    print("get_status got status_code", status.status_code)
-                    st.warning(status.content)
-                    continue
-                status = status.json()
-                print(status)
-                if status["status"] != "running" and status["status"] != "queued" :
-                    if status["msg"] != "":
-                        st.error(status["msg"])
-                    break
-                elif status["status"] == "queued":
-                    started_time = time.time()
-                    queue_length = requests.get(base_url+"/queue_length").json()
-                    progress_bar.write(f"There are {queue_length['length']} tasks in the queue")
-                elif status["progress"] == 0.0:
-                    progressed = min(0.5 * (time.time() - started_time) / 80.0, 0.5) #estimate 80s for strotts
-                    progress_bar.progress(progressed)
-                else:
-                    progress_bar.progress(min(0.5 + status["progress"] / 2.0, 1.0))
-            vp_res = requests.get(base_url+"/get_vp", params={"task_id": task_id})
-            if vp_res.status_code != 200:
-                st.warning("got status" + str(vp_res.status_code))
-                vp_res.raise_for_status()
-            else:
-                vp = np.load(BytesIO(vp_res.content))["vp"] 
-                print("received vp from server")
-                print("got numpy array", vp.shape)
-                vp = torch.from_numpy(vp).cuda()
-
-                effect, preset, param_set = get_default_settings("minimal_pipeline")
-                effect.enable_checkpoints()
-                effect.cuda()
-                content_cuda = np_to_torch(content).cuda()
-                with torch.no_grad():
-                    result_cuda = effect(content_cuda, vp)
-                img_res = Image.fromarray((torch_to_np(result_cuda) * 255.0).astype(np.uint8))
-                st.image(img_res)

tasks.py

@@ -0,0 +1,135 @@
+import base64
+import datetime
+import os
+import sys
+from io import BytesIO
+from pathlib import Path
+import numpy as np
+import requests
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import time
+import streamlit as st
+from demo_config import HUGGING_FACE, WORKER_URL
+
+
+
+PACKAGE_PARENT = 'wise'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
+
+from parameter_optimization.parametric_styletransfer import single_optimize
+from parameter_optimization.parametric_styletransfer import CONFIG as ST_CONFIG
+from parameter_optimization.strotss_org import strotss, pil_resize_long_edge_to
+from helpers import torch_to_np, np_to_torch
+
+def retrieve_for_results_from_server():
+    task_id = st.session_state['current_server_task_id']
+    vp_res = requests.get(WORKER_URL+"/get_vp", params={"task_id": task_id})
+    image_res = requests.get(WORKER_URL+"/get_image", params={"task_id": task_id})
+    if vp_res.status_code != 200 or image_res.status_code != 200:
+        st.warning("got status for " + WORKER_URL+"/get_vp" + str(vp_res.status_code))
+        st.warning("got status for " + WORKER_URL+"/image_res" + str(image_res.status_code))
+        st.session_state['current_server_task_id'] = None
+        vp_res.raise_for_status()
+        image_res.raise_for_status()
+    else:
+        st.session_state['current_server_task_id'] = None
+        vp = np.load(BytesIO(vp_res.content))["vp"] 
+        print("received vp from server")
+        print("got numpy array", vp.shape)
+        vp = torch.from_numpy(vp).cuda()
+        image = Image.open(BytesIO(image_res.content))
+        print("received image from server")
+        image = np_to_torch(np.asarray(image)).cuda()
+
+        st.session_state["effect_input"] = image
+        st.session_state["result_vp"] = vp
+
+
+def monitor_task(progress_placeholder):
+    task_id = st.session_state['current_server_task_id']
+
+    started_time = time.time()
+    retries = 3
+    while True:
+        status = requests.get(WORKER_URL+"/get_status", params={"task_id": task_id})
+        if status.status_code != 200:
+            print("get_status got status_code", status.status_code)
+            st.warning(status.content)
+            retries -= 1
+            if retries == 0:
+                return
+            else:
+                time.sleep(2)
+                continue
+        status = status.json()
+        print(status)
+        if status["status"] != "running" and status["status"] != "queued" :
+            if status["msg"] != "":
+                print("got error for task", task_id, ":", status["msg"])
+                progress_placeholder.error(status["msg"])
+                st.session_state['current_server_task_id'] = None
+                st.stop()
+            if status["status"] == "finished":
+                retrieve_for_results_from_server()
+            return
+        elif status["status"] == "queued":
+            started_time = time.time()
+            queue_length = requests.get(WORKER_URL+"/queue_length").json()
+            progress_placeholder.write(f"There are {queue_length['length']} tasks in the queue")
+        elif status["progress"] == 0.0:
+            progressed = min(0.5 * (time.time() - started_time) / 80.0, 0.5) #estimate 80s for strotts
+            progress_placeholder.progress(progressed)
+        else:
+            progress_placeholder.progress(min(0.5 + status["progress"] / 2.0, 1.0))
+
+        time.sleep(2)
+
+
+def optimize_on_server(content, style, result_image_placeholder):
+    url = WORKER_URL + "/upload"
+    content_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
+    style_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
+    asp_c, asp_s =  content.height / content.width, style.height / style.width
+    if any(a < 0.5 or a > 2.0 for a in (asp_c, asp_s)):
+        result_image_placeholder.error('aspect ratio must be <= 2')
+        st.stop()
+    content = pil_resize_long_edge_to(content, 1024)
+    content.save(content_path)
+    style = pil_resize_long_edge_to(style, 1024)
+    style.save(style_path)
+    files = {'style-image': open(style_path, "rb"), "content-image": open(content_path, "rb")}
+    print("start-optimizing")
+    task_id_res = requests.post(url, files=files)
+    if task_id_res.status_code != 200:
+        result_image_placeholder.error(task_id_res.content)
+        st.stop()
+    else:
+        task_id = task_id_res.json()['task_id']
+        st.session_state['current_server_task_id'] = task_id
+
+    monitor_task(result_image_placeholder)
+
+def optimize_params(effect, preset, content, style, result_image_placeholder):
+    result_image_placeholder.text("Executing NST to create reference image..")
+    base_dir = f"result/{datetime.datetime.now().strftime(r'%Y-%m-%d %H.%Mh %Ss')}"
+    os.makedirs(base_dir)
+    reference = strotss(pil_resize_long_edge_to(content, 1024),
+                        pil_resize_long_edge_to(style, 1024), content_weight=16.0,
+                        device=torch.device("cuda"), space="uniform")
+    progress_bar = result_image_placeholder.progress(0.0)
+    ref_save_path = os.path.join(base_dir, "reference.jpg")
+    content_save_path = os.path.join(base_dir, "content.jpg")
+    resize_to = 720
+    reference = pil_resize_long_edge_to(reference, resize_to)
+    reference.save(ref_save_path)
+    content.save(content_save_path)
+    ST_CONFIG["n_iterations"] = 300
+    
+    vp, content_img_cuda = single_optimize(effect, preset, "l1", content_save_path, str(ref_save_path),
+                                        write_video=False, base_dir=base_dir,
+                                        iter_callback=lambda i: progress_bar.progress(
+                                            float(i) / ST_CONFIG["n_iterations"]))
+    st.session_state["effect_input"], st.session_state["result_vp"]  = content_img_cuda.detach(), vp.cuda().detach()

worker/requirements.txt

@@ -3,6 +3,7 @@ imageio-ffmpeg
 scipy
 Pillow
 numpy
+matplotlib
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch
 torchvision

worker/serve.py

@@ -175,10 +175,12 @@ class StylerQueue:
     def queue_task(self, *args):
         global total_task_count
         total_task_count += 1
-        task = StyleTask(total_task_count, *args)
+        task_id = abs(hash(str(time.time())))
+        print("queued task num. ", total_task_count, "with ID", task_id)
+        task = StyleTask(task_id, *args)
         self.queued_tasks.append(task)
 
-        return total_task_count
+        return task_id
 
     def get_task(self, task_id):
         if self.running_task is not None and self.running_task.task_id == task_id:
@@ -281,4 +283,4 @@ def get_vp():
 
 
 if __name__ == '__main__':
-    app.run(debug=False, host="0.0.0.0",port=5000)
+    app.run(debug=False, host="0.0.0.0",port=8600)