init

app.py

@@ -33,6 +33,7 @@
 #       step_fn: src.diffusion.stateful_flow_matching.sampling.ode_step_fn
 import random
 import os
+import time
 import spaces
 import torch
 import argparse
@@ -69,8 +70,8 @@ def load_model(weight_dict, denoiser):
 class Pipeline:
     def __init__(self, vae, denoiser, conditioner, resolution):
         self.vae = vae
-        self.denoiser = denoiser
-        self.conditioner = conditioner
+        self.denoiser = denoiser.cuda()
+        self.conditioner = conditioner.cuda()
         self.resolution = resolution
         self.tmp_dir = tempfile.TemporaryDirectory(prefix="traj_gifs_")
         # self.denoiser.compile()
@@ -78,10 +79,10 @@ class Pipeline:
     def __del__(self):
         self.tmp_dir.cleanup()
 
-    @spaces.GPU
+    # @spaces.GPU
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-    def __call__(self, y, seed, image_height, image_width, num_steps, guidance, timeshift, order):
+    def __call__(self, y, seed, num_steps, guidance, timeshift, order):
         diffusion_sampler = AdamLMSampler(
             order=order,
             scheduler=LinearScheduler(),
@@ -91,18 +92,18 @@ class Pipeline:
             timeshift=timeshift
         )
         generator = torch.Generator(device="cpu").manual_seed(seed)
-        image_height = image_height // 32 * 32
-        image_width = image_width // 32 * 32
-        self.denoiser.decoder_patch_scaling_h = image_height / 512
-        self.denoiser.decoder_patch_scaling_w = image_width / 512
-        xT = torch.randn((1, 3, image_height, image_width), device="cpu", dtype=torch.float32,
+        xT = torch.randn((1, 3, 512, 512), device="cpu", dtype=torch.float32,
                          generator=generator).cuda()
+
+        start = time.time()
         with torch.no_grad():
             condition, uncondition = conditioner([y,]*1)
+        print("conditioner:",time.time() - start)
 
-
+        start = time.time()
         # Sample images:
         samples, trajs = diffusion_sampler(denoiser, xT, condition, uncondition, return_x_trajs=True)
+        print("diffusion:",time.time() - start)
 
         def decode_images(samples):
             samples = vae.decode(samples)
@@ -114,35 +115,35 @@ class Pipeline:
                 images.append(image)
             return images
 
-        def decode_trajs(trajs):
-            cat_trajs = torch.stack(trajs, dim=0).permute(1, 0, 2, 3, 4)
-            animations = []
-            for i in range(cat_trajs.shape[0]):
-                frames = decode_images(
-                    cat_trajs[i]
-                )
-                # 生成唯一文件名（结合seed和样本索引，避免冲突）
-                gif_filename = f"{random.randint(0, 100000)}.gif"
-                gif_path = os.path.join(self.tmp_dir.name, gif_filename)
-                frames[0].save(
-                    gif_path,
-                    format="GIF",
-                    append_images=frames[1:],
-                    save_all=True,
-                    duration=200,
-                    loop=0
-                )
-                animations.append(gif_path)
-            return animations
+        # def decode_trajs(trajs):
+        #     cat_trajs = torch.stack(trajs, dim=0).permute(1, 0, 2, 3, 4)
+        #     animations = []
+        #     for i in range(cat_trajs.shape[0]):
+        #         frames = decode_images(
+        #             cat_trajs[i]
+        #         )
+        #         # 生成唯一文件名（结合seed和样本索引，避免冲突）
+        #         gif_filename = f"{random.randint(0, 100000)}.gif"
+        #         gif_path = os.path.join(self.tmp_dir.name, gif_filename)
+        #         frames[0].save(
+        #             gif_path,
+        #             format="GIF",
+        #             append_images=frames[1:],
+        #             save_all=True,
+        #             duration=200,
+        #             loop=0
+        #         )
+        #         animations.append(gif_path)
+        #     return animations
 
         images = decode_images(samples)
-        animations = decode_trajs(trajs)
+        # animations = decode_trajs(trajs)
 
-        return images, animations
+        return images[0]
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--config", type=str, default="configs_t2i/inference_heavydecoder.yaml")
+    parser.add_argument("--config", type=str, default="configs_t2i/sft_res512.yaml")
     parser.add_argument("--resolution", type=int, default=512)
     parser.add_argument("--model_id", type=str, default="MCG-NJU/PixNerd-XXL-P16-T2I")
     parser.add_argument("--ckpt_path", type=str, default="models")
@@ -167,6 +168,7 @@ if __name__ == "__main__":
     ckpt = torch.load(ckpt_path, map_location="cpu")
     denoiser = load_model(ckpt, denoiser)
     denoiser = denoiser.cuda()
+    conditioner = conditioner.cuda()
     vae = vae.cuda()
     denoiser.eval()
 
@@ -179,27 +181,23 @@ if __name__ == "__main__":
             with gr.Column(scale=1):
                 num_steps = gr.Slider(minimum=1, maximum=100, step=1, label="num steps", value=25)
                 guidance = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, label="CFG", value=4.0)
-                image_height = gr.Slider(minimum=128, maximum=1024, step=32, label="image height", value=512)
-                image_width = gr.Slider(minimum=128, maximum=1024, step=32, label="image width", value=512)
                 label = gr.Textbox(label="positive prompt", value="a photo of a cat")
                 seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="seed", value=0)
                 timeshift = gr.Slider(minimum=0.1, maximum=5.0, step=0.1, label="timeshift", value=3.0)
                 order = gr.Slider(minimum=1, maximum=4, step=1, label="order", value=2)
             with gr.Column(scale=2):
                 btn = gr.Button("Generate")
-                output_sample = gr.Gallery(label="Images", columns=2, rows=2)
-            with gr.Column(scale=2):
-                output_trajs = gr.Gallery(label="Trajs of Diffusion", columns=2, rows=2)
+                output_sample = gr.Image(label="Images")
+            # with gr.Column(scale=2):
+            #     output_trajs = gr.Gallery(label="Trajs of Diffusion", columns=2, rows=2)
 
         btn.click(fn=pipeline,
                   inputs=[
                       label,
                       seed,
-                      image_height,
-                      image_width,
                       num_steps,
                       guidance,
                       timeshift,
                       order
-                  ], outputs=[output_sample, output_trajs])
+                  ], outputs=[output_sample])
     demo.launch()