Update app.py

app.py

@@ -1,67 +1,57 @@
+import base64
 import os
-import sys
+import time
+from io import BytesIO
+from multiprocessing import Process
 
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-import numpy as np
 import streamlit as st
 from PIL import Image
 
-import clip
-from dalle.models import Dalle
-from dalle.utils.utils import clip_score, download
+import requests
 
-url = "https://arena.kakaocdn.net/brainrepo/models/minDALL-E/57b008f02ceaa02b779c8b7463143315/1.3B.tar.gz"
-root = os.path.expanduser("~/.cache/minDALLE")
-filename = os.path.basename(url)
-pathname = filename[:-len('.tar.gz')]
 
-expected_md5 = url.split("/")[-2]
-download_target = os.path.join(root, filename)
-result_path = os.path.join(root, pathname)
+def start_server():
+    os.system("uvicorn server:app --port 8080 --host 0.0.0.0 --workers 2")
 
-if not os.path.exists(result_path):
-    result_path = download(url, root)
 
+def load_models():
+    if not is_port_in_use(8080):
+        with st.spinner(text="Loading models, please wait..."):
+            proc = Process(target=start_server, args=(), daemon=True)
+            proc.start()
+            while not is_port_in_use(8080):
+                time.sleep(1)
+            st.success("Model server started.")
+    else:
+        st.success("Model server already running...")
+    st.session_state["models_loaded"] = True
 
 
-device = "cpu"
-model = Dalle.from_pretrained(result_path)  # This will automatically download the pretrained model.
-model.to(device=device)
+def is_port_in_use(port):
+    import socket
 
-model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)
-model_clip.to(device=device)
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(("0.0.0.0", port)) == 0
 
 
-def sample(prompt):
-    # Sampling
-    images = (
-        model.sampling(prompt=prompt, top_k=256, top_p=None, softmax_temperature=1.0, num_candidates=3, device=device)
-        .cpu()
-        .numpy()
-    )
-    images = np.transpose(images, (0, 2, 3, 1))
+def generate(prompt):
+    correct_request = f"http://0.0.0.0:8080/correct?prompt={prompt}"
+    response = requests.get(correct_request)
+    images = response.json()["images"]
+    images = [Image.open(BytesIO(base64.b64decode(img))) for img in images]
+    return images
 
-    # CLIP Re-ranking
-    rank = clip_score(
-        prompt=prompt, images=images, model_clip=model_clip, preprocess_clip=preprocess_clip, device=device
-    )
 
-    # Save images
-    images = images[rank]
-    # print(rank, images.shape)
-    pil_images = []
-    for i in range(len(images)):
-        im = Image.fromarray((images[i] * 255).astype(np.uint8))
-        pil_images.append(im)
-
-    # im = Image.fromarray((images[0] * 255).astype(np.uint8))
-    return pil_images
+if "models_loaded" not in st.session_state:
+    st.session_state["models_loaded"] = False
 
 
 st.header("minDALL-E")
 st.subheader("Generate images from text")
 
+if not st.session_state["models_loaded"]:
+    load_models()
+
 prompt = st.text_input("What do you want to see?")
 
 DEBUG = False
@@ -90,9 +80,9 @@ if prompt != "":
     )
 
     print(f"Getting selections: {prompt}")
-    selected = sample(prompt)
+    selected = generate(prompt)
 
-    margin = 0.1  #for better position of zoom in arrow
+    margin = 0.1  # for better position of zoom in arrow
     n_columns = 3
     cols = st.columns([1] + [margin, 1] * (n_columns - 1))
     for i, img in enumerate(selected):