Update app.py

app.py

@@ -7,7 +7,7 @@ import openshape
 import transformers
 from PIL import Image
 from huggingface_hub import HfFolder, snapshot_download
-from demo_support import retrieval, utils, lvis
+from demo_support import retrieval, generation, utils, lvis
 from collections import OrderedDict
 
 @st.cache_resource
@@ -23,6 +23,7 @@ def load_openclip():
             clip_model.cuda()
     return clip_model, clip_prep
 
+
 @st.cache_resource
 def load_openshape(name, to_cpu=False):
     pce = openshape.load_pc_encoder(name)
@@ -30,6 +31,7 @@ def load_openshape(name, to_cpu=False):
         pce = pce.cpu()
     return pce
 
+
 def retrieval_filter_expand():
     sim_th = st.sidebar.slider("Similarity Threshold", 0.05, 0.5, 0.1, key='rsimth')
     tag = ""
@@ -47,6 +49,7 @@ def retrieval_filter_expand():
     )
     return sim_th, filter_fn
 
+
 def retrieval_results(results):
     st.caption("Click the link to view the 3D shape")
     for i in range(len(results) // 4):
@@ -65,76 +68,6 @@ def retrieval_results(results):
                 st.markdown(f"[{quote_name}]({ext_link})")
 
 
-
-def demo_captioning():
-    with st.form("capform"):
-        cond_scale = st.slider('Conditioning Scale', 0.0, 4.0, 2.0, 0.1, key='capcondscl')
-
-def demo_pc2img():
-    with st.form("sdform"):
-        prompt = st.text_input("Prompt (Optional)", key='sdtprompt')
-
-def demo_retrieval():
-    with tab_pc:
-        with st.form("rpcform"):
-            k = st.slider("Number of items to retrieve", 1, 100, 16, key='rpc')
-            load_data = utils.input_3d_shape('rpcinput')
-            sim_th, filter_fn = retrieval_filter_expand('pc')
-            if st.form_submit_button("Retrieve with Point Cloud"):
-                prog.progress(0.49, "Computing Embeddings")
-                pc = load_data(prog)
-                col2 = utils.render_pc(pc)
-                ref_dev = next(model_g14.parameters()).device
-                enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu()
-
-                sim = torch.matmul(torch.nn.functional.normalize(lvis.feats, dim=-1), torch.nn.functional.normalize(enc, dim=-1).squeeze())
-                argsort = torch.argsort(sim, descending=True)
-                pred = OrderedDict((lvis.categories[i], sim[i]) for i in argsort if i < len(lvis.categories))
-                with col2:
-                    for i, (cat, sim) in zip(range(5), pred.items()):
-                        st.text(cat)
-                        st.caption("Similarity %.4f" % sim)
-                
-                prog.progress(0.7, "Running Retrieval")
-                retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
-                
-                prog.progress(1.0, "Idle")
-
-    with tab_img:
-        with st.form("rimgform"):
-            k = st.slider("Number of items to retrieve", 1, 100, 16, key='rimage')
-            pic = st.file_uploader("Upload an Image", key='rimageinput')
-            sim_th, filter_fn = retrieval_filter_expand('image')
-            if st.form_submit_button("Retrieve with Image"):
-                prog.progress(0.49, "Computing Embeddings")
-                img = Image.open(pic)
-                st.image(img)
-                device = clip_model.device
-                tn = clip_prep(images=[img], return_tensors="pt").to(device)
-                enc = clip_model.get_image_features(pixel_values=tn['pixel_values'].type(half)).float().cpu()
-                
-                prog.progress(0.7, "Running Retrieval")
-                retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
-                
-                prog.progress(1.0, "Idle")
-
-    with tab_text:
-        with st.form("rtextform"):
-            k = st.slider("Number of items to retrieve", 1, 100, 16, key='rtext')
-            text = st.text_input("Input Text", key='rtextinput')
-            sim_th, filter_fn = retrieval_filter_expand('text')
-            if st.form_submit_button("Retrieve with Text"):
-                prog.progress(0.49, "Computing Embeddings")
-                device = clip_model.device
-                tn = clip_prep(text=[text], return_tensors='pt', truncation=True, max_length=76).to(device)
-                enc = clip_model.get_text_features(**tn).float().cpu()
-
-                prog.progress(0.7, "Running Retrieval")
-                retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
-                
-                prog.progress(1.0, "Idle")
-
-
 def classification_lvis(load_data):
     pc = load_data(prog)
     col2 = utils.render_pc(pc)
@@ -151,6 +84,7 @@ def classification_lvis(load_data):
             st.caption("Similarity %.4f" % sim)
     prog.progress(1.0, "Idle")
 
+
 def classification_custom(load_data, cats):
     pc = load_data(prog)
     col2 = utils.render_pc(pc)
@@ -174,7 +108,7 @@ def classification_custom(load_data, cats):
 
 def retrieval_pc(load_data, k, sim_th, filter_fn):
     pc = load_data(prog)
-    prog.progress(0.49, "Computing Embeddings")
+    prog.progress(0.5, "Computing Embeddings")
     col2 = utils.render_pc(pc)
     ref_dev = next(model_g14.parameters()).device
     enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu()
@@ -192,9 +126,10 @@ def retrieval_pc(load_data, k, sim_th, filter_fn):
     
     prog.progress(1.0, "Idle")
 
+
 def retrieval_img(pic, k, sim_th, filter_fn):
     img = Image.open(pic)
-    prog.progress(0.49, "Computing Embeddings")
+    prog.progress(0.5, "Computing Embeddings")
     st.image(img)
     device = clip_model.device
     tn = clip_prep(images=[img], return_tensors="pt").to(device)
@@ -205,8 +140,9 @@ def retrieval_img(pic, k, sim_th, filter_fn):
     
     prog.progress(1.0, "Idle")
 
+
 def retrieval_text(text, k, sim_th, filter_fn):
-    prog.progress(0.49, "Computing Embeddings")
+    prog.progress(0.5, "Computing Embeddings")
     device = clip_model.device
     tn = clip_prep(text=[text], return_tensors='pt', truncation=True, max_length=76).to(device)
     enc = clip_model.get_text_features(**tn).float().cpu()
@@ -216,9 +152,10 @@ def retrieval_text(text, k, sim_th, filter_fn):
     
     prog.progress(1.0, "Idle")
 
+
 def generation_img(load_data, prompt, noise_scale, cfg_scale, steps):
     pc = load_data(prog)
-    prog.progress(0.49, "Running Generation")
+    prog.progress(0.5, "Running Generation")
     col2 = utils.render_pc(pc)
     if torch.cuda.is_available():
         with sys.clip_move_lock:
@@ -226,6 +163,27 @@ def generation_img(load_data, prompt, noise_scale, cfg_scale, steps):
 
     width = 640
     height = 640
+    img = generation.pc_to_image(
+        model_g14, pc, prompt, noise_scale, width, height, cfg_scale, steps,
+        lambda i, t, _: prog.progress(0.49 + i / (steps + 1) / 2, "Running Diffusion Step %d" % i)
+    )
+    if torch.cuda.is_available():
+        with sys.clip_move_lock:
+            clip_model.cuda()
+    with col2:
+        st.image(img)
+    prog.progress(1.0, "Idle")
+
+
+def generation_text(load_data, cond_scale):
+    pc = load_data(prog)
+    prog.progress(0.5, "Running Generation")
+    col2 = utils.render_pc(pc)
+    
+    cap = generation.pc_to_text(model_g14, pc, cond_scale)
+    st.text(cap)
+    prog.progress(1.0, "Idle")
+
 
 try:
     f32 = numpy.float32
@@ -294,7 +252,7 @@ try:
         if generation_mode == "PointCloud-to-Image":
             st.title("Image Generation")
             prog = st.progress(0.0, "Idle")
-            prompt = st.sidebar.text_input("Prompt (Optional)", key='sdtprompt')
+            prompt = st.sidebar.text_input("Prompt (Optional)", key='gprompt')
             noise_scale = st.sidebar.slider('Variation Level', 0, 5, 1)
             cfg_scale = st.sidebar.slider('Guidance Scale', 0.0, 30.0, 10.0)
             steps = st.sidebar.slider('Diffusion Steps', 8, 50, 25)
@@ -303,8 +261,9 @@ try:
         elif generation_mode == "PointCloud-to-Text":
             st.title("Text Generation")
             prog = st.progress(0.0, "Idle")
+            cond_scale = st.sidebar.slider('Conditioning Scale', 0.0, 4.0, 2.0, 0.1, key='gcond')
             if st.sidebar.button("submit"):
-                pc = st.text_input("Input pc", key='rtextinput')
+                generation_text(load_data, cond_scale)
 
 except Exception:
     import traceback