Updated sampling.py

.gitignore

@@ -9,7 +9,7 @@ _exampleImages/
 _trash/
 
 1.3B.tar.gz
-
+ 
 stage1_last.ckpt
 
-stage2_last.ckpt
+stage2_last.ckpt

dalle/utils/sampling.py

@@ -68,6 +68,8 @@ def sampling(model: torch.nn.Module,
     pbar = tqdm(range(max_seq_len), total=max_seq_len) if is_tqdm else range(max_seq_len)
     pos_enc_tokens = get_positional_encoding(tokens, mode='1d')
 
+    #my_bar = st.progress(0)
+
     for cnt, h in enumerate(pbar):
         if code is None:
             code_ = None
@@ -93,8 +95,6 @@ def sampling(model: torch.nn.Module,
         else:
             past.append(present)
 
-        st.session_state.bar = cnt/max_seq_len
-
         logits = cutoff_topk_logits(logits, top_k)
         probs = F.softmax(logits, dim=-1)
         probs = cutoff_topp_probs(probs, top_p)
@@ -102,6 +102,14 @@ def sampling(model: torch.nn.Module,
         idx = torch.multinomial(probs, num_samples=1).clone().detach()
         code = idx if code is None else torch.cat([code, idx], axis=1)
 
+        #print(cnt/max_seq_len)
+        if(st.session_state.page != 0):
+            break
+
+        st.session_state.bar.progress(cnt/max_seq_len)
+
+        #my_bar.progress(cnt/max_seq_len)
+
     del past
     return code
 
@@ -151,4 +159,4 @@ def sampling_igpt(model: torch.nn.Module,
         code = idx if code is None else torch.cat([code, idx], axis=1)
 
     del past
-    return code
+    return code

minDALL-E

@@ -1 +0,0 @@
-Subproject commit e5480076b9634e9dc097e1892157ed2cf15a2f86