app demo

app.py

@@ -63,14 +63,14 @@ def load_model(weight_dict, denosier):
 
 
 class Pipeline:
-    def __init__(self, vae, denoiser, conditioner, diffusion_sampler, resolution):
+    def __init__(self, vae, denoiser, conditioner, diffusion_sampler, resolution, classlabels2ids):
         self.vae = vae
         self.denoiser = denoiser
         self.conditioner = conditioner
         self.diffusion_sampler = diffusion_sampler
         self.resolution = resolution
+        self.classlabels2ids = classlabels2ids
 
-    @spaces.GPU
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def __call__(self, y, num_images, seed, num_steps, guidance, state_refresh_rate, guidance_interval_min, guidance_interval_max, timeshift):
@@ -83,7 +83,7 @@ class Pipeline:
         generator = torch.Generator(device="cuda").manual_seed(seed)
         xT = torch.randn((num_images, 4, self.resolution//8, self.resolution//8), device="cuda", dtype=torch.float32, generator=generator)
         with torch.no_grad():
-            condition, uncondition = conditioner([y,]*num_images)
+            condition, uncondition = conditioner([self.classlabels2ids[y],]*num_images)
         # Sample images:
         samples = diffusion_sampler(denoiser, xT, condition, uncondition)
         samples = vae.decode(samples)
@@ -136,7 +136,15 @@ if __name__ == "__main__":
     vae = vae.cuda()
     denoiser.eval()
 
-    pipeline = Pipeline(vae, denoiser, conditioner, diffusion_sampler, args.resolution)
+    # read imagenet classlabels
+    with open("imagenet_classlabels.txt", "r") as f:
+        classlabels = f.readlines()
+        classlabels = [label.strip() for label in classlabels]
+
+    classlabels2id = {label: i for i, label in enumerate(classlabels)}
+    id2classlabels = {i: label for i, label in enumerate(classlabels)}
+
+    pipeline = Pipeline(vae, denoiser, conditioner, diffusion_sampler, args.resolution, classlabels2id)
 
     with gr.Blocks() as demo:
         gr.Markdown("DDT")
@@ -144,12 +152,14 @@ if __name__ == "__main__":
             with gr.Column(scale=1):
                 num_steps = gr.Slider(minimum=1, maximum=100, step=1, label="num steps", value=50)
                 guidance = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, label="CFG", value=4.0)
-                num_images = gr.Slider(minimum=1, maximum=10, step=1, label="num images", value=8)
-                label = gr.Slider(minimum=0, maximum=999, step=1, label="label", value=948)
+                num_images = gr.Slider(minimum=1, maximum=10, step=1, label="num images", value=4)
+                label = gr.Dropdown(choices=classlabels, value=id2classlabels[948], label="label")
                 seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="seed", value=0)
                 state_refresh_rate = gr.Slider(minimum=1, maximum=10, step=1, label="encoder reuse", value=1)
-                guidance_interval_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="interval guidance min", value=0.0)
-                guidance_interval_max = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="interval guidance max", value=1.0)
+                guidance_interval_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="interval guidance min",
+                                                  value=0.0)
+                guidance_interval_max = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="interval guidance max",
+                                                  value=1.0)
                 timeshift = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, label="timeshift", value=1.0)
             with gr.Column(scale=2):
                 btn = gr.Button("Generate")
@@ -167,4 +177,4 @@ if __name__ == "__main__":
                       guidance_interval_max,
                       timeshift
                   ], outputs=[output])
-    demo.launch(server_name="0.0.0.0", server_port=7861)
+    demo.launch()

src/diffusion/stateful_flow_matching/sharing_sampling.py

@@ -109,7 +109,7 @@ class EulerSampler(BaseSampler):
         timesteps.reverse()
 
         print("recompute timesteps solved by DP: ", timesteps)
-        return timesteps[:-1]
+        return timesteps[:-1][:self.num_recompute_timesteps]
 
     def _impl_sampling(self, net, noise, condition, uncondition):
         """