Update utils.py

utils.py

@@ -66,7 +66,6 @@ def layout(*args):
     st.markdown(str(foot), unsafe_allow_html=True)
 
 
-
 def footer():
     myargs = [
         "Created by ",
@@ -96,7 +95,6 @@ def footer():
     height=600,
     )
 
-
 model = False
 def generate(prompt,crazy,k):
     global model
@@ -113,7 +111,11 @@ def generate(prompt,crazy,k):
     set_seed(np.random.randint(0,10000))
 
     # Sampling
-    images = model.sampling(prompt=prompt,
+    newPrompt = prompt
+    if("architecture" not in prompt.lower() ):
+        newPrompt += " architecture"
+
+    images = model.sampling(prompt=newPrompt,
                             top_k=2048,
                             top_p=None,
                             softmax_temperature=crazy,
@@ -124,7 +126,7 @@ def generate(prompt,crazy,k):
     # CLIP Re-ranking
     model_clip, preprocess_clip = clip.load("ViT-B/32", device=device)
     model_clip.to(device=device)
-    rank = clip_score(prompt=prompt,
+    rank = clip_score(prompt=newPrompt,
                       images=images,
                       model_clip=model_clip,
                       preprocess_clip=preprocess_clip,
@@ -143,35 +145,37 @@ def generate(prompt,crazy,k):
 
 def drawGrid():
     master = {}
-    order = 0
-
-    #print(st.session_state.results)
     
     for r in st.session_state.results[::-1]:
         _txt = r['prompt']+" "+str(r['crazy'])+" "+str(r['k'])
-
         if(_txt not in master):
             master[_txt] = [r]
-            order += 1
         else:
             master[_txt].append(r)
 
-    
-    for m in master:
-        #with placeholder.container():
+
+    for i in st.session_state.images:
+        im = st.empty()
+
+
+    placeholder = st.empty()
+    with placeholder.container():
+        
+        for m in master:
             
-        txt = master[m][0]['prompt']+" (temperature:"+ str(master[m][0]['crazy']) + ", top k:" + str(master[m][0]['k']) + ")"
-        st.subheader(txt)
-        col1, col2, col3 = st.columns(3)   
-
-        for ix, item in enumerate(master[m]):
-            if ix % 3 == 0:
-                with col1:
-                    st.image(item["image"])
-            if ix % 3 == 1:
-                with col2:
-                    st.image(item["image"])
-            if ix % 3 == 2:
-                with col3:
-                    st.image(item["image"])
+            txt = master[m][0]['prompt']+" (temperature:"+ str(master[m][0]['crazy']) + ", top k:" + str(master[m][0]['k']) + ")"
+            st.subheader(txt)
+            col1, col2, col3 = st.columns(3)  
+
+            for ix, item in enumerate(master[m]):
+                if ix % 3 == 0: 
+                    with col1:
+                        st.session_state.images.append(st.image(item["image"]))  
+                if ix % 3 == 1:
+                    with col2:
+                        st.session_state.images.append(st.image(item["image"]))
+                if ix % 3 == 2:
+                    with col3:
+                        st.session_state.images.append(st.image(item["image"]))
+