feat: progress bat

app.py

@@ -70,43 +70,43 @@ encoder = get_encoder()
 model = get_model()
 diffusion = get_diffusion()
 
-sample_fn = diffusion.ddim_sample_loop
-
 st.title("Lang2mol-Diff")
 text_input = st.text_area("Enter molecule description")
-with st.spinner("Please wait ..."):
-    output = tokenizer(
-        text_input,
-        max_length=256,
-        truncation=True,
-        padding="max_length",
-        add_special_tokens=True,
-        return_tensors="pt",
-        return_attention_mask=True,
-    )
-    caption_state = encoder(
-        input_ids=output["input_ids"],
-        attention_mask=output["attention_mask"],
-    ).last_hidden_state
-    caption_mask = output["attention_mask"]
+button = st.button("Submit")
+if button:
+    with st.spinner("Please wait..."):
+        output = tokenizer(
+            text_input,
+            max_length=256,
+            truncation=True,
+            padding="max_length",
+            add_special_tokens=True,
+            return_tensors="pt",
+            return_attention_mask=True,
+        )
+        caption_state = encoder(
+            input_ids=output["input_ids"],
+            attention_mask=output["attention_mask"],
+        ).last_hidden_state
+        caption_mask = output["attention_mask"]
 
-    outputs = sample_fn(
-        model,
-        (1, 256, 32),
-        clip_denoised=False,
-        denoised_fn=None,
-        model_kwargs={},
-        top_p=1.0,
-        progress=True,
-        caption=(caption_state, caption_mask),
-    )
-    logits = model.get_logits(torch.tensor(outputs))
-    cands = torch.topk(logits, k=1, dim=-1)
-    outputs = cands.indices
-    outputs = outputs.squeeze(-1)
-    outputs = tokenizer.decode(outputs)
-    result = sf.decoder(
-        outputs[0].replace("<pad>", "").replace("</s>", "").replace("\t", "")
-    ).replace("\t", "")
+        outputs = diffusion.p_sample_loop(
+            model,
+            (1, 256, 32),
+            clip_denoised=False,
+            denoised_fn=None,
+            model_kwargs={},
+            top_p=1.0,
+            progress=True,
+            caption=(caption_state, caption_mask),
+        )
+        logits = model.get_logits(torch.tensor(outputs))
+        cands = torch.topk(logits, k=1, dim=-1)
+        outputs = cands.indices
+        outputs = outputs.squeeze(-1)
+        outputs = tokenizer.decode(outputs)
+        result = sf.decoder(
+            outputs[0].replace("<pad>", "").replace("</s>", "").replace("\t", "")
+        ).replace("\t", "")
 
-    st.write(result)
+        st.write(result)

src/improved_diffusion/gaussian_diffusion.py

@@ -9,7 +9,7 @@ import enum
 import math
 import torch
 import numpy as np
-
+import streamlit as st
 from .nn import mean_flat
 from .losses import normal_kl, discretized_gaussian_log_likelihood
 
@@ -667,16 +667,19 @@ class GaussianDiffusion:
         # print(indices[-10:])
         if progress:
             # Lazy import so that we don't depend on tqdm.
-            from tqdm.auto import tqdm
+            # from tqdm.auto import tqdm
+
+            # indices = tqdm(indices)
+            pass
 
-            indices = tqdm(indices)
         if caption is not None:
             print("Text Guiding Generation ......")
             caption = (
                 caption[0].to(img.device),
                 caption[1].to(img.device),
             )  # (caption_state, caption_mask)
-        for i in indices:
+        my_bar = st.progress(0, text="Processing")
+        for pro, i in enumerate(indices):
             t = torch.tensor([i] * shape[0], device=device)
             with torch.no_grad():
                 out = self.p_sample(
@@ -691,6 +694,8 @@ class GaussianDiffusion:
                 )
                 yield out
                 img = out["sample"]
+            my_bar.progress(pro + 1, text="Processing")
+        my_bar.empty()
 
     def p_sample_loop_langevin_progressive(
         self,