Merge pull request #55 from storytold/22-april-kasisnu

22 april kasisnu

tts/tacotron2_v1_early_fakeyou/Dockerfile

@@ -106,7 +106,7 @@ WORKDIR models/tts
 # NB: Setuptools fix: https://github.com/tensorflow/tensorflow/issues/34302#issuecomment-554450289
 RUN python3.6 -m pip install --upgrade pip setuptools \
   && python3.6 -m pip install -r requirements-tacotron-python36.txt \
-  && python3.6 -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 
+  && python3.6 -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 && python3.6 -m pip cache purge
 
 # Final image
 #  FROM ubuntu:xenial
@@ -140,7 +140,7 @@ RUN apt-get install -y libespeak-ng1
 RUN python3 --version
 
 # TODO: Use venv to do this instead.
-RUN pip3 install gdown
+RUN pip3 install gdown && pip3 cache purge
 
 # NB(echelon,2022-07-14): Run a newer version of pytorch which will support RTX3090.
 #  Fixes bug, "NVIDIA GeForce RTX 3080 with CUDA capability sm_86 is not compatible with the current PyTorch installation."

workflows/preview_server/install.py

@@ -4,6 +4,10 @@ import json
 import os
 from pathlib import Path
 import requests
+from multiprocessing import Pool
+
+
+
 
 # set cwd to script location
 os.chdir(os.path.dirname(os.path.realpath(__file__)))
@@ -54,33 +58,39 @@ def run_comfyui_setup():
     run_command(f"{venv_python} initialize_comfy_extensions.py")
 
 
+
+def install_dependency(dep):
+    location = dep['location']
+    url = dep['url']
+
+    # Convert the location to a Path object and create parent directories if they don't exist
+    path = Path(location)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if path.is_file():
+        print(f"{path} is already installed")
+        return
+
+    # Download the file
+    response = requests.get(url)
+    if response.status_code == 200:
+        # Write the content to the file in the specified location
+        with open(path, 'wb') as file:
+            file.write(response.content)
+        print(f"Installed {path}")
+    else:
+        print(f"Failed to download {url}")
+
+
+
 def install_dependencies(json_file_path='dependencies.json'):
     # Load the dependency JSON
     with open(json_file_path, 'r') as file:
         dependencies = json.load(file)
 
-    for dep in dependencies:
-        # Extract location and URL
-        location = dep['location']
-        url = dep['url']
 
-        # Convert the location to a Path object and create parent directories if they don't exist
-        path = Path(location)
-        path.parent.mkdir(parents=True, exist_ok=True)
-
-        if path.is_file():
-            print(f"{path} is already installed")
-            continue
-
-        # Download the file
-        response = requests.get(url)
-        if response.status_code == 200:
-            # Write the content to the file in the specified location
-            with open(path, 'wb') as file:
-                file.write(response.content)
-            print(f"Installed {path}")
-        else:
-            print(f"Failed to download {url}")
+    with Pool(5) as pool:
+        pool.starmap(install_dependency, [(dep,) for dep in dependencies])
 
 
 def main():

workflows/preview_server/preview_server.py

@@ -12,28 +12,69 @@ import random
 from typing import Dict, Any, Optional
 import mimetypes
 
-from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Request
-from fastapi.responses import FileResponse
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Request, Response
 import requests
 from jsonpath_ng import jsonpath, parse
 from moviepy.editor import ImageClip
+from fastapi.middleware.cors import CORSMiddleware
+import urllib
+import websockets
 
+# CORS header for which websites can import this content in a frame.
+# This *technically* shouldn't be required, but the XHR requests are 
+# being made from a controlled frame.
+#
+# https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Security-Policy/frame-ancestors
+#
+# NB(bt,2024-03-30): The *.netlify.app is dangerous.
+CORS_WHITELIST_ORIGINS = [
+    "https://studio.storyteller.ai",
+    "https://storytellerstudio.netlify.app",
+    "https://branch-name--storytellerstudio.netlify.app",
+    "https://deploy-preview-123--storytellerstudio.netlify.app",
+    "https://pipeline-gottagofast.netlify.app",
+    "https://branch-whatever--pipeline-gottagofast.netlify.app",
+    "https://fakeyou.com",
+    "https://*.fakeyou.com",
+    "https://storyteller.ai",
+    "https://*.storyteller.ai",
+    "https://*.netlify.app",
+    "http://*.fakeyou.com:7000",
+    "http://*.fakeyou.com:7001",
+    "http://*.fakeyou.com:7002",
+    "http://*.storyteller.ai:7000",
+    "http://*.storyteller.ai:7001",
+    "http://*.storyteller.ai:7002",
+]
+
+
+# TODO(kasisnu): Remove ugly globals
 app = FastAPI()
-lock = asyncio.Lock()  # Create an asyncio lock
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=CORS_WHITELIST_ORIGINS,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 SERVER_IP = os.environ.get("SERVER_IP", "127.0.0.1:8188")
 TIMEOUT_SECONDS = os.environ.get("TIMEOUT_SECONDS", 1000)
 PROMPT_ENDPOINT = f"http://{SERVER_IP}/prompt"
 V2V_WORKFLOWS_DIRECTORY = os.environ.get("V2V_WORKFLOWS_DIRECTORY",
-                                         "/home/salt/clone/storyteller/V2V-StylesList/V2VWorkflows/")
+                                         "/home/kasisnu/go/src/github.com/storytold/V2V-StylesList/V2VWorkflows/")
 COMFY_ROOT_DIRECTORY = os.environ.get("COMFY_ROOT_DIRECTORY",
-                                      "/home/salt/clone/storyteller/storyteller-ml/workflows/preview_server/ComfyUI")
+                                      "/home/kasisnu/go/src/github.com/comfyanonymous/ComfyUI")
 SERVER_START_TIMEOUT = int(os.environ.get("SERVER_START_TIMEOUT", 60))
 STYLE_LOCK = os.environ.get("STYLE_LOCK", None)
-PREVIEW_WORKFLOW_NAME = os.environ.get("PREVIEW_WORKFLOW_NAME", None)
+PREVIEW_WORKFLOW_NAME = os.environ.get("PREVIEW_WORKFLOW_NAME", 'workflow_api_instantpreview.json')
 MAPPING_NAME = os.environ.get("MAPPING_NAME", "mapping_two.json")
 AUTH_URL = os.environ.get("AUTH_URL", "https://api.fakeyou.com/v1/session")
 
+# websocket.enableTrace(True)
+client_id = str(uuid.uuid4())
+
+
 style_to_filename = {
     "anime_2_5d": "1_2.5d_anime_model.json",
     "anime_2d_flat": "2_2d_flat_anime_model.json",
@@ -55,30 +96,6 @@ style_to_filename = {
 
 allowed_keys = ["style", "positive_prompt", "negative_prompt"]
 
-# CORS header for which websites can import this content in a frame.
-# This *technically* shouldn't be required, but the XHR requests are 
-# being made from a controlled frame.
-#
-# https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Security-Policy/frame-ancestors
-#
-# NB(bt,2024-03-30): The *.netlify.app is dangerous.
-CORS_FRAME_ANCESTORS = [
-    "frame-ancestors 'self'",
-    "https://studio.storyteller.ai",
-    "https://storytellerstudio.netlify.app",
-    "https://fakeyou.com",
-    "https://*.fakeyou.com",
-    "https://storyteller.ai",
-    "https://*.storyteller.ai",
-    "https://*.netlify.app",
-    "http://*.fakeyou.com:7000",
-    "http://*.fakeyou.com:7001",
-    "http://*.fakeyou.com:7002",
-    "http://*.storyteller.ai:7000",
-    "http://*.storyteller.ai:7001",
-    "http://*.storyteller.ai:7002",
-    ";",
-]
 
 
 def can_use_vst(cookies):
@@ -148,6 +165,7 @@ def signal_handler(signum, frame):
     exit(0)
 
 
+
 @app.post("/preview/")
 async def preview_request(request_obj: Request, request: str = Form(...), input_file: UploadFile = File(...)):
     cookies = request_obj.cookies
@@ -159,53 +177,42 @@ async def preview_request(request_obj: Request, request: str = Form(...), input_
     if STYLE_LOCK:
         if request["style"] != STYLE_LOCK:
             raise HTTPException(status_code=400, detail=f"This endpoint is for '{STYLE_LOCK}' style only")
+    try:
+        request_style = request["style"]
+        style_filename = style_to_filename[request["style"]]
+        positive_prompt = request.get("positive_prompt", None)
+        negative_prompt = request.get("negative_prompt", None)
+        mime_type = input_file.content_type
+        input_filename_ext = ''
+        if 'video' in mime_type:
+            input_filename_ext = ".mp4"
+        elif 'image' in mime_type:
+            input_filename_ext = ".jpeg"
+        else:
+            raise HTTPException(status_code=400, detail="Unsupported file type")
 
-    async with lock:
-        try:
-            style_filename = style_to_filename[request["style"]]
-            positive_prompt = request.get("positive_prompt", None)
-            negative_prompt = request.get("negative_prompt", None)
+        temp_filename = f"temp_files/{uuid.uuid4()}{input_filename_ext}"
+        Path(temp_filename).parent.mkdir(parents=True, exist_ok=True)
+        with open(temp_filename, "wb") as buffer:
+            shutil.copyfileobj(input_file.file, buffer)
+        if 'image' in mime_type:
+            clip = ImageClip(temp_filename)
+            fps = 24
+            clip_duration = 1 / fps
+            clip = clip.set_duration(clip_duration)
+            input_filename = f"temp_files/{uuid.uuid4()}.mp4"
+            clip.write_videofile(input_filename, fps=fps)
 
-            temp_filename = f"temp_files/{uuid.uuid4()}"
-            Path(temp_filename).parent.mkdir(parents=True, exist_ok=True)
-            with open(temp_filename, "wb") as buffer:
-                shutil.copyfileobj(input_file.file, buffer)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to save file: {str(e)}")
 
-            mime_type, _ = mimetypes.guess_type(input_file.filename)
+    # Process video (or image converted to video) and generate preview
+    preview_image_path = await process_video(request_style, input_filename, style_filename, positive_prompt, negative_prompt)
 
-            if 'video' in mime_type:
-                input_filename = f"{temp_filename}.mp4"
-            elif 'image' in mime_type:
-                input_filename = f"{temp_filename}.{mime_type.split('/')[-1]}"
-            else:
-                raise HTTPException(status_code=400, detail="Unsupported file type")
-
-            Path(temp_filename).rename(input_filename)
-
-            if 'image' in mime_type:
-                clip = ImageClip(input_filename)
-                fps = 24
-                clip_duration = 1 / fps
-                clip = clip.set_duration(clip_duration)
-                input_filename = f"temp_files/{uuid.uuid4()}.mp4"
-                clip.write_videofile(input_filename, fps=fps)
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to save file: {str(e)}")
-
-        # Process video (or image converted to video) and generate preview
-        preview_image_path = await process_video(input_filename, style_filename, positive_prompt, negative_prompt)
-
-    csp_frame_ancestors_value = ' '.join(CORS_FRAME_ANCESTORS[:-1])
-    headers = {
-        "Access-Control-Allow-Origin": "*",  # Be cautious with this setting in production environments.
-        "Content-Security-Policy": f"frame-ancestors {csp_frame_ancestors_value}",
-    }
-
-    return FileResponse(preview_image_path, media_type='image/jpeg', headers=headers)
+    return Response(preview_image_path, media_type='image/jpeg')
 
 
-async def process_video(video_path: str, style_filename, positive_prompt=None, negative_prompt=None) -> Path:
+async def process_video(request_style: str, video_path: str, style_filename, positive_prompt=None, negative_prompt=None) -> Path:
     # generate random seed (otherwise video will cache completely and not generate a new output)
     styles_directory = Path(V2V_WORKFLOWS_DIRECTORY) / "styles"
     workflow_directory = Path(V2V_WORKFLOWS_DIRECTORY) / "workflows"
@@ -213,7 +220,6 @@ async def process_video(video_path: str, style_filename, positive_prompt=None, n
 
     style_json = json.load(open(styles_directory / style_filename))
     workflow_json = json.load(open(workflow_directory / PREVIEW_WORKFLOW_NAME))
-    # mapping_json = json.load(open(mappings_directory / style_json["mapping_name"]))
     mapping_json = json.load(open(mappings_directory / MAPPING_NAME))
 
     json_mods = style_json["modifications"]
@@ -229,54 +235,23 @@ async def process_video(video_path: str, style_filename, positive_prompt=None, n
         pass
 
     input_video_path = Path(COMFY_ROOT_DIRECTORY) / "input/input.mp4"
-    preview_image_path = Path(COMFY_ROOT_DIRECTORY) / "output/vid2vid/SparseImagePreview_00001_.png"
     shutil.move(video_path, input_video_path)
-    if os.path.exists(preview_image_path):
-        os.remove(preview_image_path)
-
-    response = requests.get(PROMPT_ENDPOINT, timeout=5)
-    if response.status_code == 200:
-        print("Server is healthy")
-    else:
-        print("Server is unhealthy")
-        print(f"Status code: {response.status_code}")
-        print(f"Response: {response.text}")
-        raise HTTPException(status_code=500, detail="Server is unhealthy")
-
-    # make sure there are no jobs already
-    queue_json = response.json()
-    if queue_json["exec_info"]["queue_remaining"]:
-        print("There are jobs in the queue already. Exiting...")
-        raise HTTPException(status_code=500, detail="There are jobs in the queue already (this shouldn't happen)")
-    else:
-        print("Queue is empty")
-
-    p = {"prompt": json_prompt}
+    p = {"prompt": json_prompt, "client_id": client_id, "extra_data": {"style": request_style}}
     data = json.dumps(p).encode('utf-8')
+    execution_start_time = time.perf_counter()
     response = requests.post(PROMPT_ENDPOINT, data=data)
     if response.status_code == 200:
         print("Job enqueued successfully")
     else:
-        print("Job failed to queue")
-        print(f"Status code: {response.status_code}")
-        print(f"Response: {response.text}")
         raise HTTPException(status_code=500, detail=f"Failed to queue job")
+    prompt_id = response.json()["prompt_id"]
+    async with websockets.connect("ws://{}/ws?clientId={}".format(SERVER_IP, client_id)) as ws:
+        images = await get_images(ws, prompt_id, '1517')
+        current_time = time.perf_counter()
+        execution_time = current_time - execution_start_time
+        print('request executed in {:.2f} seconds'.format(execution_time))
+        return images['1517'][0]
 
-    start_time = time.time()
-    while True:
-        response = requests.get(PROMPT_ENDPOINT, timeout=5)
-        queue_json = response.json()
-        if queue_json["exec_info"]["queue_remaining"]:
-            pass
-        else:
-            print("Job is done")
-            break
-        if time.time() - start_time > TIMEOUT_SECONDS:
-            print("Timeout reached and job didn't finish!")
-            raise HTTPException(status_code=500, detail="Timeout reached and job didn't finish!")
-        time.sleep(1)
-
-    return preview_image_path
 
 
 def get_jsonpath_mods(style_mods: Dict[str, Any], mapping_json: Dict[str, Any], pos_in: Optional[str] = None,
@@ -346,17 +321,57 @@ def apply_jsonpath_mods(jsonpath_mods: Dict[str, Any], workflow_json: Dict[str,
     return workflow_json
 
 
+def get_image(filename, subfolder, folder_type):
+    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    url_values = urllib.parse.urlencode(data)
+    with urllib.request.urlopen("http://{}/view?{}".format(SERVER_IP, url_values)) as response:
+        return response.read()
+
+
+
+def get_history(prompt_id):
+    with urllib.request.urlopen("http://{}/history/{}".format(SERVER_IP, prompt_id)) as response:
+        return json.loads(response.read())
+
+async def get_images(ws, prompt_id, filter_node_id=None):
+    output_images = {}
+    while True:
+        out = await ws.recv()
+        if isinstance(out, str):
+            message = json.loads(out)
+            if message['type'] == 'executing':
+                data = message['data']
+                if data['node'] is None and data['prompt_id'] == prompt_id:
+                    break #Execution is done
+        else:
+            continue #previews are binary data
+
+    history = get_history(prompt_id)[prompt_id]
+    for o in history['outputs']:
+        for node_id in history['outputs']:
+            if filter_node_id is not None and node_id != filter_node_id:
+                continue
+            node_output = history['outputs'][node_id]
+            images_output = []
+            if 'images' in node_output:
+                for image in node_output['images']:
+                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
+                    images_output.append(image_data)
+            output_images[node_id] = images_output
+
+    return output_images
+
+
+
 if __name__ == "__main__":
     import uvicorn
-    import asyncio
-
     # install ComfyUI
-    os.system("python3 install.py")
+    # os.system("python3 install.py")
 
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)
+    comfyui_process = subprocess.Popen(["venv/bin/python", "main.py", "--highvram"], cwd=COMFY_ROOT_DIRECTORY)
 
-    comfyui_process = subprocess.Popen(["venv/bin/python", "main.py"], cwd=COMFY_ROOT_DIRECTORY)
 
     # block until comfyUI server is up
     start_time = time.time()
@@ -377,4 +392,4 @@ if __name__ == "__main__":
     print("ComfyUI server running, starting FastAPI server")
     uvicorn.run(app, host="0.0.0.0", port=8000)
     # cleanup
-    terminate_process(comfyui_process)
+    terminate_process(comfyui_process)

workflows/preview_server/server_requirements.txt

@@ -4,4 +4,5 @@ fastapi~=0.110.0
 python-multipart~=0.0.9
 pydantic~=2.6.4
 jsonpath-ng~=1.6.1
-moviepy
+moviepy
+websockets

workflows/preview_server/test_request.py

@@ -1,23 +1,62 @@
 import json
-
+import time
 import requests
 
-# url = 'http://207.189.112.61:31605/preview'
-# url = "http://207.189.112.61:31605/preview"
-url = "http://localhost:8000/preview"
-payload = {"style": "paper_origami", "positive_prompt": "test pos in", "negative_prompt": "test neg in"}
-files = {'input_file': open('test_image.jpg', 'rb')}
+# TEST_INPUT_FILE = 'test_image.jpg'
+TEST_INPUT_FILE = 'test_image_10pc.jpg'
 
-cookies = {
-    "session": "<put your session cookie here>"
+# url = 'http://b3104:8005/preview/'
+# url = "http://207.189.112.61:31605/preview"
+# url = 'http://localhost:8000/preview/'
+
+url = 'https://funnel.tailce84f.ts.net/preview/'
+style_to_filename = {
+    "anime_2_5d": "1_2.5d_anime_model.json",
+    "anime_2d_flat": "2_2d_flat_anime_model.json",
+    "cartoon_3d": "3_3d_cartoon_style.json",
+    "comic_book": "4_comic_book_model.json",
+    "anime_ghibli": "5_ghibli_anime_model.json",
+    "ink_punk": "6_ink_punk.json",
+    "ink_splash": "7_ink_splash.json",
+    "ink_bw_style": "8_ink_w_and_b_style.json",
+    "jojo_style": "9_jojo_style.json",
+    "paper_origami": "10_paper_origami.json",
+    # "pixel_art": "11_pixel_art.json",
+    "pop_art": "12_pop_art.json",
+    "realistic_1": "13_realistic_1.json",
+    "realistic_2": "14_realistic_2.json",
+    "anime_retro_neon": "15_retro_neon_anime_90.json",
+    "anime_standard": "16_standard_anime_model.json",
 }
 
-response = requests.post(url, data={'request': json.dumps(payload)}, files=files, cookies=cookies)
-if response.status_code == 200:
-    # Success
-    with open('preview_image.jpeg', 'wb') as f:
-        f.write(response.content)
-else:
-    # Handle request error
-    print(f"Request failed: {response.status_code}")
-    print(response.text)
+all_styles = list(style_to_filename.keys())
+
+import random
+random.shuffle(all_styles)
+
+styles = list(all_styles)[:6]
+# random.shuffle(styles)
+
+test_style_list = styles + styles
+# for style, filename in list(style_to_filename.items())[:2]:
+for style in all_styles:
+    for i in range(3):
+        # print(f"Testing style: {style}")
+        payload = {"style": style, "positive_prompt": "test pos in", "negative_prompt": "test neg in"}
+        with open(TEST_INPUT_FILE , 'rb') as input_file:
+            # files = {'input_file': ('input_file', input_file, 'image/jpeg')}
+            files = {'input_file':  input_file}
+            execution_start_time = time.perf_counter()
+            with requests.post(url, data={'request': json.dumps(payload)}, files=files) as response:
+                if response.status_code == 200:
+                    # write output file suffixed with timestamp
+                    timestamp = time.strftime('%Y%m%d%H%M%S')
+                    with open(f'outputs/preview_image_{style}_{timestamp}.jpeg', 'wb') as f:
+                        f.write(response.content)
+                else:
+                    # Handle request error
+                    print(f"Request failed: {response.status_code}")
+                    print(response.text)
+            current_time = time.perf_counter()
+            execution_time = current_time - execution_start_time
+            print('{} executed in {:.2f} seconds'.format(style, execution_time))