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()