Merge pull request #1 from MaxReimann/server_test

.gitignore

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

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,33 +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.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()
@@ -223,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://ava.hpi3d.de:8600"

docker-compose.yml

@@ -0,0 +1,11 @@
+version: "2.3"
+
+services:
+  worker:
+    build:
+        context: ./
+        dockerfile: ./project/server/Dockerfile_worker
+    image: wise-worker
+    container_name: wise-eccv-optim-worker
+    ports:
+      - 8600:8600

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/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"] 

worker/requirements.txt

@@ -0,0 +1,11 @@
+imageio
+imageio-ffmpeg
+scipy
+Pillow
+numpy
+matplotlib
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
+torchvision
+Flask
+Flask-Reuploaded

worker/serve.py

@@ -0,0 +1,286 @@
+import datetime
+import os
+from pathlib import Path
+import sys
+from flask import Flask, jsonify, request, send_file, abort
+from flask_uploads import UploadSet, configure_uploads, IMAGES
+from werkzeug.exceptions import default_exceptions
+from werkzeug.exceptions import HTTPException, NotFound 
+import json
+import torch
+import time
+import threading
+import traceback
+from PIL import Image
+import numpy as np
+
+PACKAGE_PARENT = '..'
+WISE_DIR = '../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)))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, WISE_DIR)))
+
+
+
+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
+from effects import get_default_settings, MinimalPipelineEffect 
+
+class JSONExceptionHandler(object):
+
+    def __init__(self, app=None):
+        if app:
+            self.init_app(app)
+
+    def std_handler(self, error):
+        response = jsonify(message=error.message)
+        response.status_code = error.code if isinstance(error, HTTPException) else 500
+        return response
+
+
+    def init_app(self, app):
+        self.app = app
+        self.register(HTTPException)
+        for code, v in default_exceptions.items():
+            self.register(code)
+
+    def register(self, exception_or_code, handler=None):
+        self.app.errorhandler(exception_or_code)(handler or self.std_handler)
+
+
+
+app = Flask(__name__)
+handler = JSONExceptionHandler(app)
+
+image_folder = 'img_received'
+photos = UploadSet('photos', IMAGES)
+app.config['UPLOADED_PHOTOS_DEST'] = image_folder
+configure_uploads(app, photos)
+
+class Args(object):
+    def __init__(self, initial_data):
+        for key in initial_data:
+            setattr(self, key, initial_data[key])
+    def set_attributes(self, val_dict):
+        for key in val_dict:
+            setattr(self, key, val_dict[key])
+
+default_args = {
+        "output_image" : "output.jpg",
+        ## values always set by request ##
+        "content_image": "",
+        "style_image": "",
+        "output_vp": "",
+        "iters": 500
+}
+
+
+total_task_count = 0
+
+class NeuralOptimizer():
+    def __init__(self, args) -> None:
+        self.cur_iteration = 0
+        self.args = args
+
+    def optimize(self):
+        base_dir = f"result/{datetime.datetime.now().strftime(r'%Y-%m-%d %H.%Mh %Ss')}"
+        os.makedirs(base_dir)
+
+        content = Image.open(self.args.content_image)
+        style = Image.open(self.args.style_image)
+
+        def set_iter(iter):
+            self.cur_iteration=iter
+
+        effect, preset, _ = get_default_settings("minimal_pipeline")
+        effect.enable_checkpoints()
+
+        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")
+
+        ref_save_path = os.path.join(base_dir, "reference.jpg")
+        resize_to = 720
+        reference = pil_resize_long_edge_to(reference, resize_to)
+        reference.save(ref_save_path)
+        ST_CONFIG["n_iterations"] = self.args.iters
+        vp, content_img_cuda = single_optimize(effect, preset, "l1", self.args.content_image, str(ref_save_path),
+                                            write_video=False, base_dir=base_dir,
+                                            iter_callback=set_iter)
+
+        output = Image.fromarray(torch_to_np(content_img_cuda.detach().cpu() * 255.0).astype(np.uint8))
+        output.save(self.args.output_image)
+        # torch.save (vp.detach().clone(), self.args.output_vp)
+        # preset_tensor = effect.vpd.preset_tensor(preset, np_to_torch(np.array(content)).cuda(), add_local_dims=True)
+        np.savez_compressed(self.args.output_vp, vp=vp.detach().cpu().numpy())
+
+        
+
+class StyleTask:
+    def __init__(self, task_id, style_filename, content_filename):
+        self.content_filename=content_filename
+        self.style_filename=style_filename
+
+        self.status = "queued"
+        self.task_id = task_id
+        self.error_msg = ""
+        self.output_filename = content_filename.split(".")[0] + "_output.jpg"
+        self.vp_output_filename = content_filename.split(".")[0] + "_output.npz"
+
+        # global neural_optimizer
+        # if neural_optimizer is None:
+        #     neural_optimizer = NeuralOptimizer(Args(default_args))
+
+        self.neural_optimizer = NeuralOptimizer(Args(default_args))
+    
+    def start(self):
+        self.neural_optimizer.args.set_attributes(default_args)
+
+        self.neural_optimizer.args.style_image = os.path.join(image_folder, self.style_filename)
+        self.neural_optimizer.args.content_image = os.path.join(image_folder, self.content_filename)
+        self.neural_optimizer.args.output_image = os.path.join(image_folder, self.output_filename)
+        self.neural_optimizer.args.output_vp = os.path.join(image_folder, self.vp_output_filename)
+ 
+        thread = threading.Thread(target=self.run, args=())
+        thread.daemon = True                            # Daemonize thread
+        thread.start()                                  # Start the execution
+
+    def run(self):
+        self.status = "running"
+        try:
+            self.neural_optimizer.optimize()
+        except Exception as e:
+            print("Error in task %d :"%(self.task_id), str(e))
+            traceback.print_exc()
+
+            self.status = "error"
+            self.error_msg = str(e)
+            return
+
+        self.status = "finished"
+        print("finished styling task: " + str(self.task_id))
+
+class StylerQueue:
+    queued_tasks = []
+    finished_tasks = []
+    running_task = None
+
+    def __init__(self):
+        thread = threading.Thread(target=self.status_checker, args=())
+        thread.daemon = True                            # Daemonize thread
+        thread.start()                                  # Start the execution
+
+    def queue_task(self, *args):
+        global total_task_count
+        total_task_count += 1
+        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 task_id
+
+    def get_task(self, task_id):
+        if self.running_task is not None and self.running_task.task_id == task_id:
+            return self.running_task
+        task = next((task for task in self.queued_tasks + self.finished_tasks if task.task_id == task_id), None)
+        return task
+
+    def status_checker(self):
+        while True:
+            time.sleep(0.3)
+
+            if self.running_task is None:
+                if len(self.queued_tasks) > 0:
+                    self.running_task = self.queued_tasks[0]
+                    self.running_task.start()
+                    self.queued_tasks = self.queued_tasks[1:]
+            elif self.running_task.status == "finished" or self.running_task.status == "error": 
+                self.finished_tasks.append(self.running_task)
+                if len(self.queued_tasks) > 0:
+                    self.running_task = self.queued_tasks[0]
+                    self.running_task.start()
+                    self.queued_tasks = self.queued_tasks[1:]
+                else:
+                    self.running_task = None
+
+styler_queue = StylerQueue()
+
+
+@app.route('/upload', methods=['POST'])
+def upload():
+    if 'style-image' in request.files and \
+        'content-image' in request.files:
+
+        style_filename = photos.save(request.files['style-image'])
+        content_filename = photos.save(request.files['content-image'])
+
+        job_id = styler_queue.queue_task(style_filename, content_filename)
+        print('added new stylization task', style_filename, content_filename)
+
+        return jsonify({"task_id": job_id})
+    abort(jsonify(message="request needs style, content image"), 400)
+
+@app.route('/get_status')
+def get_status():
+    task_id = int(request.args.get("task_id"))
+    task = styler_queue.get_task(task_id)
+
+    if task is None:
+        abort(jsonify(message="task with id %d not found"%task_id), 400) 
+
+    status = {
+        "status": task.status,
+        "msg": task.error_msg
+    }
+
+    if task.status == "running":
+        if isinstance(task, StyleTask):
+            status["progress"] = float(task.neural_optimizer.cur_iteration) / float(default_args["iters"]) 
+
+    return jsonify(status)
+
+@app.route('/queue_length')
+def get_queue_length():
+    tasks = len(styler_queue.queued_tasks)
+    if styler_queue.running_task is not None:
+        tasks += 1
+
+    status = {
+        "length": tasks
+    }
+
+    return jsonify(status)
+
+
+@app.route('/get_image')
+def get_image():
+    task_id = int(request.args.get("task_id"))
+    task = styler_queue.get_task(task_id)
+
+    if task is None:
+        abort(jsonify(message="task with id %d not found"%task_id), 400)
+
+    if task.status != "finished":
+        abort(jsonify(message="task with id %d not in finished state"%task_id), 400)
+
+    return send_file(os.path.join(image_folder, task.output_filename), mimetype='image/jpg')
+
+@app.route('/get_vp')
+def get_vp():
+    task_id = int(request.args.get("task_id"))
+    task = styler_queue.get_task(task_id)
+
+    if task is None:
+        abort(jsonify(message="task with id %d not found"%task_id), 400)
+
+    if task.status != "finished":
+        abort(jsonify(message="task with id %d not in finished state"%task_id), 400)
+
+    return send_file(os.path.join(image_folder, task.vp_output_filename), mimetype='application/zip')
+
+
+if __name__ == '__main__':
+    app.run(debug=False, host="0.0.0.0",port=8600)