t

app.py

@@ -122,17 +122,36 @@ ab_t[0] = 1
 # construct model
 nn_model = ContextUnet(in_channels=3, n_feat=n_feat, n_cfeat=n_cfeat, height=height).to(device)
 
+# sample quickly using DDIM
+@torch.no_grad()
+def sample_ddim(n_sample, n=20):
+    # x_T ~ N(0, 1), sample initial noise
+    samples = torch.randn(n_sample, 3, height, height).to(device)
+
+    # array to keep track of generated steps for plotting
+    intermediate = []
+    step_size = timesteps // n
+    for i in range(timesteps, 0, -step_size):
+        print(f'sampling timestep {i:3d}', end='\r')
+
+        # reshape time tensor
+        t = torch.tensor([i / timesteps])[:, None, None, None].to(device)
+
+        eps = nn_model(samples, t)    # predict noise e_(x_t,t)
+        samples = denoise_ddim(samples, i, i - step_size, eps)
+        intermediate.append(samples.detach().cpu().numpy())
+
+    intermediate = np.stack(intermediate)
+    return samples, intermediate
+
 def greet(input):
-	prompt = f"""
-Recommend complementary shop combinations which match well with the shop(s) described in the following text, which is delimited by triple backticks. Rank by synergy: \
-Text: ```{input}```
-    """
-	response = prompt
+	samples, intermediate = sample_ddim(32, n=25)
+	response = intermediate[-1]
 	return response
 
 #iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 #iface.launch()
 
 #iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Text to find entities", lines=2)], outputs=[gr.HighlightedText(label="Text with entities")], title="NER with dslim/bert-base-NER", description="Find entities using the `dslim/bert-base-NER` model under the hood!", allow_flagging="never", examples=["My name is Andrew and I live in California", "My name is Poli and work at HuggingFace"])
-iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Co-Retailing Business")], outputs="text")
+iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Co-Retailing Business")], outputs="image")
 iface.launch()