Merge pull request #162 from borisdayma/demo-improvements

app/streamlit/app.py

@@ -1,28 +1,10 @@
 #!/usr/bin/env python
 # coding: utf-8
 
-import base64
-from io import BytesIO
+from datetime import datetime
 
-import requests
 import streamlit as st
-from PIL import Image
-
-
-class ServiceError(Exception):
-    def __init__(self, status_code):
-        self.status_code = status_code
-
-
-def get_images_from_backend(prompt, backend_url):
-    r = requests.post(backend_url, json={"prompt": prompt})
-    if r.status_code == 200:
-        images = r.json()["images"]
-        images = [Image.open(BytesIO(base64.b64decode(img))) for img in images]
-        return images
-    else:
-        raise ServiceError(r.status_code)
-
+from backend import ServiceError, get_images_from_backend, get_model_version
 
 st.sidebar.markdown(
     """
@@ -45,7 +27,7 @@ DALL·E mini is an AI model that generates images from any prompt you give!
 </p>
 
 <p style='text-align: center'>
-Created by Boris Dayma et al. 2021
+Created by Boris Dayma et al. 2021-2022
 <br/>
 <a href="https://github.com/borisdayma/dalle-mini" target="_blank">GitHub</a> | <a href="https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA" target="_blank">Project Report</a>
 </p>
@@ -84,8 +66,8 @@ if prompt != "":
     )
 
     try:
-        backend_url = st.secrets["BACKEND_SERVER"]
-        print(f"Getting selections: {prompt}")
+        backend_url = st.secrets["BACKEND_SERVER"] + "/generate"
+        print(f"{datetime.now()} Getting selections: {prompt}")
         selected = get_images_from_backend(prompt, backend_url)
 
         margin = 0.1  # for better position of zoom in arrow
@@ -95,6 +77,18 @@ if prompt != "":
             cols[(i % n_columns) * 2].image(img)
         container.markdown(f"**{prompt}**")
 
+        version_url = st.secrets["BACKEND_SERVER"] + "/version"
+        version = get_model_version(version_url)
+        st.sidebar.markdown(
+            f"<small><center>{version}</center></small>", unsafe_allow_html=True
+        )
+
+        st.markdown(
+            f"""
+        These results have been obtained using model `{version}` from [an ongoing training run](https://wandb.ai/dalle-mini/dalle-mini/runs/mheh9e55).
+        """
+        )
+
         st.button("Again!", key="again_button")
 
     except ServiceError as error:

app/streamlit/backend.py

@@ -0,0 +1,31 @@
+# Client requests to Dalle-Mini Backend server
+
+import base64
+from io import BytesIO
+
+import requests
+from PIL import Image
+
+
+class ServiceError(Exception):
+    def __init__(self, status_code):
+        self.status_code = status_code
+
+
+def get_images_from_backend(prompt, backend_url):
+    r = requests.post(backend_url, json={"prompt": prompt})
+    if r.status_code == 200:
+        images = r.json()["images"]
+        images = [Image.open(BytesIO(base64.b64decode(img))) for img in images]
+        return images
+    else:
+        raise ServiceError(r.status_code)
+
+
+def get_model_version(url):
+    r = requests.get(url)
+    if r.status_code == 200:
+        version = r.json()["version"]
+        return version
+    else:
+        raise ServiceError(r.status_code)