Merged super demo.

README.md

@@ -1,5 +1,5 @@
 ---
-title: OpenShape Classification Demo
+title: OpenShape Demo
 emoji: 🌖
 colorFrom: red
 colorTo: purple

app.py

@@ -1,30 +1,67 @@
+import sys
 import streamlit as st
 from huggingface_hub import HfFolder, snapshot_download
-HfFolder().save_token(st.secrets['etoken'])
-snapshot_download("OpenShape/openshape-demo-support", local_dir='.')
+
+
+@st.cache_data
+def load_support():
+    HfFolder().save_token(st.secrets['etoken'])
+    sys.path.append(snapshot_download("OpenShape/openshape-demo-support"))
+
+
+# load_support()
 
 
 import numpy
+import torch
 import openshape
-from openshape.demo import misc_utils, classification
-
+import transformers
+from PIL import Image
 
 @st.cache_resource
 def load_openshape(name):
     return openshape.load_pc_encoder(name)
 
 
+@st.cache_resource
+def load_openclip():
+    return transformers.CLIPModel.from_pretrained(
+        "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k",
+        low_cpu_mem_usage=True, torch_dtype=half,
+        offload_state_dict=True
+    ), transformers.CLIPProcessor.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
+
+
 f32 = numpy.float32
-# clip_model, clip_prep = load_openclip()
-model_g14 = openshape.load_pc_encoder('openshape-pointbert-vitg14-rgb')
+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').cpu()
+model_l14 = load_openshape('openshape-pointbert-vitl14-rgb')
+model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
+torch.set_grad_enabled(False)
 
+from openshape.demo import misc_utils, classification, caption, sd_pc2img, retrieval
 
 st.title("OpenShape Demo")
-load_data = misc_utils.input_3d_shape()
 prog = st.progress(0.0, "Idle")
+tab_cls, tab_text, tab_img, tab_pc, tab_sd, tab_cap = st.tabs([
+    "Classification",
+    "Retrieval from Text",
+    "Retrieval from Image",
+    "Retrieval from 3D Shape",
+    "Image Generation",
+    "Captioning",
+])
 
 
-try:
+def demo_classification():
+    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:
+        st.error('Maximum 64 custom categories supported in the demo')
+        return
     if st.button("Run Classification on LVIS Categories"):
         pc = load_data(prog)
         col2 = misc_utils.render_pc(pc)
@@ -35,5 +72,127 @@ try:
                 st.text(cat)
                 st.caption("Similarity %.4f" % sim)
         prog.progress(1.0, "Idle")
-except Exception as exc:
-    st.error(repr(exc))
+    if st.button("Run Classification on Custom Categories"):
+        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).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")
+
+
+def demo_captioning():
+    load_data = misc_utils.input_3d_shape('cap')
+    cond_scale = st.slider('Conditioning Scale', 0.0, 4.0, 2.0)
+    if st.button("Generate a Caption"):
+        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")
+
+
+def demo_pc2img():
+    load_data = misc_utils.input_3d_shape('sd')
+    prompt = st.text_input("Prompt (Optional)")
+    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.button("Generate"):
+        pc = load_data(prog)
+        col2 = misc_utils.render_pc(pc)
+        prog.progress(0.49, "Running Generation")
+        if torch.cuda.is_available():
+            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():
+            clip_model.cuda()
+        with col2:
+            st.image(img)
+        prog.progress(1.0, "Idle")
+
+
+def retrieval_results(results):
+    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_text:
+        k = st.slider("# Shapes to Retrieve", 1, 100, 16, key='rtext')
+        text = st.text_input("Input Text")
+        if st.button("Run 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))
+            prog.progress(1.0, "Idle")
+
+    with tab_img:
+        k = st.slider("# Shapes to Retrieve", 1, 100, 16, key='rimage')
+        pic = st.file_uploader("Upload an Image")
+        if st.button("Run with Image"):
+            img = Image.open(pic)
+            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))
+            prog.progress(1.0, "Idle")
+
+    with tab_pc:
+        k = st.slider("# Shapes to Retrieve", 1, 100, 16, key='rpc')
+        load_data = misc_utils.input_3d_shape('retpc')
+        if st.button("Run with Shape"):
+            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))
+            prog.progress(1.0, "Idle")
+
+
+try:
+    if torch.cuda.is_available():
+        clip_model.cuda()
+    with tab_cls:
+        demo_classification()
+    with tab_cap:
+        demo_captioning()
+    with tab_sd:
+        demo_pc2img()
+    demo_retrieval()
+except Exception:
+    import traceback
+    st.error(traceback.format_exc().replace("\n", "  \n"))