Update app.py

app.py

@@ -1,33 +1,13 @@
 import sys
 import threading
 import streamlit as st
-from huggingface_hub import HfFolder, snapshot_download
-
-
-@st.cache_data
-def load_support():
-    if st.secrets.has_key('etoken'):
-        HfFolder().save_token(st.secrets['etoken'])
-    sys.path.append(snapshot_download("OpenShape/openshape-demo-support"))
-
-
-# st.set_page_config(layout='wide')
-load_support()
-
-
 import numpy
 import torch
 import openshape
 import transformers
 from PIL import Image
-
-@st.cache_resource
-def load_openshape(name, to_cpu=False):
-    pce = openshape.load_pc_encoder(name)
-    if to_cpu:
-        pce = pce.cpu()
-    return pce
-
+from huggingface_hub import HfFolder, snapshot_download
+from demo_support import retrieval
 
 @st.cache_resource
 def load_openclip():
@@ -43,95 +23,48 @@ def load_openclip():
     return clip_model, clip_prep
 
 
-f32 = numpy.float32
-half = torch.float16 if torch.cuda.is_available() else torch.bfloat16
-# clip_model, clip_prep = None, None
-clip_model, clip_prep = load_openclip()
-model_b32 = load_openshape('openshape-pointbert-vitb32-rgb', True)
-model_l14 = load_openshape('openshape-pointbert-vitl14-rgb')
-model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
-torch.set_grad_enabled(False)
-for kc, vc in st.session_state.get('state_queue', []):
-    st.session_state[kc] = vc
-st.session_state.state_queue = []
-
-
-import samples_index
-from openshape.demo import misc_utils, classification, caption, sd_pc2img, retrieval
-
-
-st.title("OpenShape Demo")
-st.caption("For faster inference without waiting in queue, you may clone the space and run it yourself.")
-prog = st.progress(0.0, "Idle")
-tab_cls, tab_img, tab_text, tab_pc, tab_sd, tab_cap = st.tabs([
-    "Classification",
-    "Retrieval w/ Image",
-    "Retrieval w/ Text",
-    "Retrieval w/ 3D",
-    "Image Generation",
-    "Captioning",
-])
-
-
-def sq(kc, vc):
-    st.session_state.state_queue.append((kc, vc))
-
-
-def reset_3d_shape_input(key):
-    # this is not working due to streamlit problems, don't use it
-    model_key = key + "_model"
-    npy_key = key + "_npy"
-    swap_key = key + "_swap"
-    sq(model_key, None)
-    sq(npy_key, None)
-    sq(swap_key, "Y is up (for most Objaverse shapes)")
-
-
-def auto_submit(key):
-    if st.session_state.get(key):
-        st.session_state[key] = False
-        return True
-    return False
-
-
-def queue_auto_submit(key):
-    st.session_state[key] = True
-    st.experimental_rerun()
-
-
-img_example_counter = 0
+def retrieval_filter_expand(key):
+    with st.expander("Filters"):
+        sim_th = st.slider("Similarity Threshold", 0.05, 0.5, 0.1, key=key + 'rtsimth')
+        tag = st.text_input("Has Tag", "", key=key + 'rthastag')
+        col1, col2 = st.columns(2)
+        face_min = int(col1.text_input("Face Count Min", "0", key=key + 'rtfcmin'))
+        face_max = int(col2.text_input("Face Count Max", "34985808", key=key + 'rtfcmax'))
+        col1, col2 = st.columns(2)
+        anim_min = int(col1.text_input("Animation Count Min", "0", key=key + 'rtacmin'))
+        anim_max = int(col2.text_input("Animation Count Max", "563", key=key + 'rtacmax'))
+        tag_n = not bool(tag.strip())
+        anim_n = not (anim_min > 0 or anim_max < 563)
+        face_n = not (face_min > 0 or face_max < 34985808)
+        filter_fn = lambda x: (
+            (anim_n or anim_min <= x['anims'] <= anim_max)
+            and (face_n or face_min <= x['faces'] <= face_max)
+            and (tag_n or tag in x['tags'])
+        )
+        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):
+        cols = st.columns(4)
+        for j in range(4):
+            idx = i * 4 + j
+            if idx >= len(results):
+                continue
+            entry = results[idx]
+            with cols[j]:
+                ext_link = f"https://objaverse.allenai.org/explore/?query={entry['u']}"
+                st.image(entry['img'])
+                # st.markdown(f"[![thumbnail {entry['desc'].replace('\n', ' ')}]({entry['img']})]({ext_link})")
+                # st.text(entry['name'])
+                quote_name = entry['name'].replace('[', '\\[').replace(']', '\\]').replace('\n', ' ')
+                st.markdown(f"[{quote_name}]({ext_link})")
 
-def image_examples(samples, ncols, return_key=None, example_text="Examples"):
-    global img_example_counter
-    trigger = False
-    with st.expander(example_text, True):
-        for i in range(len(samples) // ncols):
-            cols = st.columns(ncols)
-            for j in range(ncols):
-                idx = i * ncols + j
-                if idx >= len(samples):
-                    continue
-                entry = samples[idx]
-                with cols[j]:
-                    st.image(entry['dispi'])
-                    img_example_counter += 1
-                    with st.columns(5)[2]:
-                        this_trigger = st.button('\+', key='imgexuse%d' % img_example_counter)
-                    trigger = trigger or this_trigger
-                    if this_trigger:
-                        if return_key is None:
-                            for k, v in entry.items():
-                                if not k.startswith('disp'):
-                                    sq(k, v)
-                        else:
-                            trigger = entry[return_key]
-    return trigger
 
 
 def demo_classification():
     with st.form("clsform"):
-        load_data = misc_utils.input_3d_shape('cls')
+        #load_data = misc_utils.input_3d_shape('cls')
         cats = st.text_input("Custom Categories (64 max, separated with comma)")
         cats = [a.strip() for a in cats.split(',')]
         if len(cats) > 64:
@@ -139,185 +72,64 @@ def demo_classification():
             return
         lvis_run = st.form_submit_button("Run Classification on LVIS Categories")
         custom_run = st.form_submit_button("Run Classification on Custom Categories")
-        if lvis_run or auto_submit("clsauto"):
-            pc = load_data(prog)
-            col2 = misc_utils.render_pc(pc)
-            prog.progress(0.5, "Running Classification")
-            pred = classification.pred_lvis_sims(model_g14, pc)
-            with col2:
-                for i, (cat, sim) in zip(range(5), pred.items()):
-                    st.text(cat)
-                    st.caption("Similarity %.4f" % sim)
-            prog.progress(1.0, "Idle")
-        if custom_run:
-            pc = load_data(prog)
-            col2 = misc_utils.render_pc(pc)
-            prog.progress(0.5, "Computing Category Embeddings")
-            device = clip_model.device
-            tn = clip_prep(text=cats, return_tensors='pt', truncation=True, max_length=76, padding=True).to(device)
-            feats = clip_model.get_text_features(**tn).float().cpu()
-            prog.progress(0.5, "Running Classification")
-            pred = classification.pred_custom_sims(model_g14, pc, cats, feats)
-            with col2:
-                for i, (cat, sim) in zip(range(5), pred.items()):
-                    st.text(cat)
-                    st.caption("Similarity %.4f" % sim)
-            prog.progress(1.0, "Idle")
-    if image_examples(samples_index.classification, 3, example_text="Examples (Choose one of the following 3D shapes)"):
-        queue_auto_submit("clsauto")
-
 
 def demo_captioning():
     with st.form("capform"):
-        load_data = misc_utils.input_3d_shape('cap')
         cond_scale = st.slider('Conditioning Scale', 0.0, 4.0, 2.0, 0.1, key='capcondscl')
-        if st.form_submit_button("Generate a Caption") or auto_submit("capauto"):
-            pc = load_data(prog)
-            col2 = misc_utils.render_pc(pc)
-            prog.progress(0.5, "Running Generation")
-            cap = caption.pc_caption(model_b32, pc, cond_scale)
-            st.text(cap)
-            prog.progress(1.0, "Idle")
-    if image_examples(samples_index.cap, 3, example_text="Examples (Choose one of the following 3D shapes)"):
-        queue_auto_submit("capauto")
-
 
 def demo_pc2img():
     with st.form("sdform"):
-        load_data = misc_utils.input_3d_shape('sd')
         prompt = st.text_input("Prompt (Optional)", key='sdtprompt')
-        noise_scale = st.slider('Variation Level', 0, 5, 1)
-        cfg_scale = st.slider('Guidance Scale', 0.0, 30.0, 10.0)
-        steps = st.slider('Diffusion Steps', 8, 50, 25)
-        width = 640  # st.slider('Width', 480, 640, step=32)
-        height = 640  # st.slider('Height', 480, 640, step=32)
-        if st.form_submit_button("Generate") or auto_submit("sdauto"):
-            pc = load_data(prog)
-            col2 = misc_utils.render_pc(pc)
-            prog.progress(0.49, "Running Generation")
-            if torch.cuda.is_available():
-                with sys.clip_move_lock:
-                    clip_model.cpu()
-            img = sd_pc2img.pc_to_image(
-                model_l14, 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")
-    if image_examples(samples_index.sd, 3, example_text="Examples (Choose one of the following 3D shapes)"):
-        queue_auto_submit("sdauto")
-
-
-def retrieval_results(results):
-    st.caption("Click the link to view the 3D shape")
-    for i in range(len(results) // 4):
-        cols = st.columns(4)
-        for j in range(4):
-            idx = i * 4 + j
-            if idx >= len(results):
-                continue
-            entry = results[idx]
-            with cols[j]:
-                ext_link = f"https://objaverse.allenai.org/explore/?query={entry['u']}"
-                st.image(entry['img'])
-                # st.markdown(f"[![thumbnail {entry['desc'].replace('\n', ' ')}]({entry['img']})]({ext_link})")
-                # st.text(entry['name'])
-                quote_name = entry['name'].replace('[', '\\[').replace(']', '\\]').replace('\n', ' ')
-                st.markdown(f"[{quote_name}]({ext_link})")
 
+def demo_retrieval():
+    with tab_pc:
+        with st.form("rpcform"):
+            k = st.slider("Number of items to retrieve", 1, 100, 16, key='rpc')
+            pc = utils.load_3D_shape('rpcinput')
+            if st.form_submit_button("Retrieve with Point Cloud"):
+                prog.progress(0.49, "Computing Embeddings")
 
-def retrieval_filter_expand(key):
-    with st.expander("Filters"):
-        sim_th = st.slider("Similarity Threshold", 0.05, 0.5, 0.1, key=key + 'rtsimth')
-        tag = st.text_input("Has Tag", "", key=key + 'rthastag')
-        col1, col2 = st.columns(2)
-        face_min = int(col1.text_input("Face Count Min", "0", key=key + 'rtfcmin'))
-        face_max = int(col2.text_input("Face Count Max", "34985808", key=key + 'rtfcmax'))
-        col1, col2 = st.columns(2)
-        anim_min = int(col1.text_input("Animation Count Min", "0", key=key + 'rtacmin'))
-        anim_max = int(col2.text_input("Animation Count Max", "563", key=key + 'rtacmax'))
-        tag_n = not bool(tag.strip())
-        anim_n = not (anim_min > 0 or anim_max < 563)
-        face_n = not (face_min > 0 or face_max < 34985808)
-        filter_fn = lambda x: (
-            (anim_n or anim_min <= x['anims'] <= anim_max)
-            and (face_n or face_min <= x['faces'] <= face_max)
-            and (tag_n or tag in x['tags'])
-        )
-        return sim_th, filter_fn
 
+    with tab_img:
+        with st.form("rimgform"):
+            k = st.slider("Number of items to retrieve", 1, 100, 16, key='rimage')
+            img = st.file_uploader("Upload an Image", key='rimageinput')
+            if st.form_submit_button("Retrieve with Image"):
+                prog.progress(0.49, "Computing Embeddings")
 
-def demo_retrieval():
     with tab_text:
         with st.form("rtextform"):
-            k = st.slider("Shapes to Retrieve", 1, 100, 16, key='rtext')
-            text = st.text_input("Input Text", key="inputrtext")
+            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("Run with Text") or auto_submit("rtextauto"):
+            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)
+                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")
-        picked_sample = st.selectbox("Examples", ["Select..."] + samples_index.retrieval_texts)
-        text_last_example = st.session_state.get('text_last_example', None)
-        if text_last_example is None:
-            st.session_state.text_last_example = picked_sample
-        elif text_last_example != picked_sample and picked_sample != "Select...":
-            st.session_state.text_last_example = picked_sample
-            sq("inputrtext", picked_sample)
-            queue_auto_submit("rtextauto")
 
-    with tab_img:
-        submit = False
-        with st.form("rimgform"):
-            k = st.slider("Shapes 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("Run with Image"):
-                submit = True
-            results_container = st.container()
-        sample_got = image_examples(samples_index.iret, 4, 'rimageinput')
-        if sample_got:
-            pic = sample_got
-        if sample_got or submit:
-            img = Image.open(pic)
-            with results_container:
-                st.image(img)
-                prog.progress(0.49, "Computing Embeddings")
-                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_pc:
-        with st.form("rpcform"):
-            k = st.slider("Shapes to Retrieve", 1, 100, 16, key='rpc')
-            load_data = misc_utils.input_3d_shape('retpc')
-            sim_th, filter_fn = retrieval_filter_expand('pc')
-            if st.form_submit_button("Run with Shape") or auto_submit('rpcauto'):
-                pc = load_data(prog)
-                col2 = misc_utils.render_pc(pc)
-                prog.progress(0.49, "Computing Embeddings")
-                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()
-                prog.progress(0.7, "Running Retrieval")
-                retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
-                prog.progress(1.0, "Idle")
-        if image_examples(samples_index.pret, 3):
-            queue_auto_submit("rpcauto")
+st.title("TripletMix Demo")
+st.caption("For faster inference without waiting in queue, you may clone the space and run it yourself.")
+prog = st.progress(0.0, "Idle")
+tab_cls, tab_pc, tab_img, tab_text, tab_sd, tab_cap = st.tabs([
+    "Classification",
+    "Retrieval w/ 3D",
+    "Retrieval w/ Image",
+    "Retrieval w/ Text",
+    "Image Generation",
+    "Captioning",
+])
 
 
+f32 = numpy.float32
+half = torch.float16 if torch.cuda.is_available() else torch.bfloat16
+clip_model, clip_prep = load_openclip()
+
 try:
     with tab_cls:
         demo_classification()