Refine gradio

Dockerfile

@@ -25,4 +25,3 @@ RUN ./environment_setup.sh sana
 
 # COPY server.py server.py
 CMD ["conda", "run", "-n", "sana", "--no-capture-output", "python", "-u", "-W", "ignore", "app/app_sana.py", "--config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml", "--model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth"]
-

app/app_sana.py

@@ -19,6 +19,8 @@
 import argparse
 import os
 import random
+import socket
+import sqlite3
 import time
 import uuid
 from datetime import datetime
@@ -41,6 +43,7 @@
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 DEMO_PORT = int(os.getenv("DEMO_PORT", "15432"))
 os.environ["GRADIO_EXAMPLES_CACHE"] = "./.gradio/cache"
+COUNTER_DB = os.getenv("COUNTER_DB", ".count.db")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -109,36 +112,37 @@
 SCHEDULE_NAME = ["Flow_DPM_Solver"]
 DEFAULT_SCHEDULE_NAME = "Flow_DPM_Solver"
 NUM_IMAGES_PER_PROMPT = 1
-TEST_TIMES = 0
 INFER_SPEED = 0
-FILENAME = f"output/port{DEMO_PORT}_inference_count.txt"
 
 
-def read_inference_count():
-    global TEST_TIMES
-    try:
-        with open(FILENAME) as f:
-            count = int(f.read().strip())
-    except FileNotFoundError:
-        count = 0
-    TEST_TIMES = count
+def open_db():
+    db = sqlite3.connect(COUNTER_DB)
+    db.execute("CREATE TABLE IF NOT EXISTS counter(app CHARS PRIMARY KEY UNIQUE, value INTEGER)")
+    db.execute('INSERT OR IGNORE INTO counter(app, value) VALUES("Sana", 0)')
+    return db
+
 
-    return count
+def read_inference_count():
+    with open_db() as db:
+        cur = db.execute('SELECT value FROM counter WHERE app="Sana"')
+        db.commit()
+    return cur.fetchone()[0]
 
 
 def write_inference_count(count):
-    with open(FILENAME, "w") as f:
-        f.write(str(count))
+    count = max(0, int(count))
+    with open_db() as db:
+        db.execute(f'UPDATE counter SET value=value+{count} WHERE app="Sana"')
+        db.commit()
 
 
 def run_inference(num_imgs=1):
-    TEST_TIMES = read_inference_count()
-    TEST_TIMES += int(num_imgs)
-    write_inference_count(TEST_TIMES)
+    write_inference_count(num_imgs)
+    count = read_inference_count()
 
     return (
         f"<span style='font-size: 16px; font-weight: bold;'>Total inference runs: </span><span style='font-size: "
-        f"16px; color:red; font-weight: bold;'>{TEST_TIMES}</span>"
+        f"16px; color:red; font-weight: bold;'>{count}</span>"
     )
 
 
@@ -238,12 +242,12 @@ def generate(
     flow_dpms_inference_steps: int = 20,
     randomize_seed: bool = False,
 ):
-    global TEST_TIMES
     global INFER_SPEED
     # seed = 823753551
+    box = run_inference(num_imgs)
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
-    print(f"PORT: {DEMO_PORT}, model_path: {model_path}, time_times: {TEST_TIMES}")
+    print(f"PORT: {DEMO_PORT}, model_path: {model_path}")
     if safety_check.is_dangerous(safety_checker_tokenizer, safety_checker_model, prompt, threshold=0.2):
         prompt = "A red heart."
 
@@ -294,11 +298,11 @@ def generate(
         img,
         seed,
         f"<span style='font-size: 16px; font-weight: bold;'>Inference Speed: {INFER_SPEED:.3f} s/Img</span>",
+        box,
     )
 
 
-TEST_TIMES = read_inference_count()
-model_size = "1.6" if "D20" in args.model_path else "0.6"
+model_size = "1.6" if "1600M" in args.model_path else "0.6"
 title = f"""
     <div style='display: flex; align-items: center; justify-content: center; text-align: center;'>
         <img src="https://raw.githubusercontent.com/NVlabs/Sana/refs/heads/main/asset/logo.png" width="50%" alt="logo"/>
@@ -308,7 +312,7 @@ def generate(
         <p><span style="font-size: 36px; font-weight: bold;">Sana-{model_size}B</span><span style="font-size: 20px; font-weight: bold;">{args.image_size}px</span></p>
         <p style="font-size: 16px; font-weight: bold;">Sana: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformer</p>
         <p><span style="font-size: 16px;"><a href="https://arxiv.org/abs/2410.10629">[Paper]</a></span> <span style="font-size: 16px;"><a href="https://github.com/NVlabs/Sana">[Github(coming soon)]</a></span> <span style="font-size: 16px;"><a href="https://nvlabs.github.io/Sana">[Project]</a></span</p>
-        <p style="font-size: 16px; font-weight: bold;">Powered by <a href="https://hanlab.mit.edu/projects/dc-ae">DC-AE</a> with 32x latent space</p>, running on A6000 node.
+        <p style="font-size: 16px; font-weight: bold;">Powered by <a href="https://hanlab.mit.edu/projects/dc-ae">DC-AE</a> with 32x latent space, </p>running on node {socket.gethostname()}.
         <p style="font-size: 16px; font-weight: bold;">Unsafe word will give you a 'Red Heart' in the image instead.</p>
         """
 if model_size == "0.6":
@@ -334,9 +338,9 @@ def generate(
 .gradio-container{max-width: 640px !important}
 h1{text-align:center}
 """
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, title="Sana") as demo:
     gr.Markdown(title)
-    gr.Markdown(DESCRIPTION)
+    gr.HTML(DESCRIPTION)
     gr.DuplicateButton(
         value="Duplicate Space for private use",
         elem_id="duplicate-button",
@@ -442,8 +446,6 @@ def generate(
                     value=1,
                 )
 
-    run_button.click(fn=run_inference, inputs=num_imgs, outputs=info_box)
-
     gr.Examples(
         examples=examples,
         inputs=prompt,
@@ -480,9 +482,9 @@ def generate(
             flow_dpms_inference_steps,
             randomize_seed,
         ],
-        outputs=[result, seed, speed_box],
+        outputs=[result, seed, speed_box, info_box],
         api_name="run",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=DEMO_PORT, debug=True, share=True)
+    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=DEMO_PORT, debug=False, share=False)

app/sana_pipeline.py

@@ -231,86 +231,85 @@ def forward(
                 ),
                 torch.tensor([[1.0]], device=self.device).repeat(num_images_per_prompt, 1),
             )
+
             for _ in range(num_images_per_prompt):
-                with torch.no_grad():
-                    prompts.append(
-                        prepare_prompt_ar(prompt, self.base_ratios, device=self.device, show=False)[0].strip()
-                    )
+                prompts.append(prepare_prompt_ar(prompt, self.base_ratios, device=self.device, show=False)[0].strip())
 
-                    # prepare text feature
-                    if not self.config.text_encoder.chi_prompt:
-                        max_length_all = self.config.text_encoder.model_max_length
-                        prompts_all = prompts
-                    else:
-                        chi_prompt = "\n".join(self.config.text_encoder.chi_prompt)
-                        prompts_all = [chi_prompt + prompt for prompt in prompts]
-                        num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt))
-                        max_length_all = (
-                            num_chi_prompt_tokens + self.config.text_encoder.model_max_length - 2
-                        )  # magic number 2: [bos], [_]
-
-                    caption_token = self.tokenizer(
-                        prompts_all,
-                        max_length=max_length_all,
-                        padding="max_length",
-                        truncation=True,
-                        return_tensors="pt",
-                    ).to(device=self.device)
-                    select_index = [0] + list(range(-self.config.text_encoder.model_max_length + 1, 0))
-                    caption_embs = self.text_encoder(caption_token.input_ids, caption_token.attention_mask)[0][:, None][
-                        :, :, select_index
-                    ].to(self.weight_dtype)
-                    emb_masks = caption_token.attention_mask[:, select_index]
-                    null_y = self.null_caption_embs.repeat(len(prompts), 1, 1)[:, None].to(self.weight_dtype)
-
-                    n = len(prompts)
-                    if latents is None:
-                        z = torch.randn(
-                            n,
-                            self.config.vae.vae_latent_dim,
-                            self.latent_size_h,
-                            self.latent_size_w,
-                            generator=generator,
-                            device=self.device,
-                            dtype=self.weight_dtype,
-                        )
-                    else:
-                        z = latents.to(self.weight_dtype).to(self.device)
-                    model_kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": ar}, mask=emb_masks)
-                    if self.vis_sampler == "flow_euler":
-                        flow_solver = FlowEuler(
-                            self.model,
-                            condition=caption_embs,
-                            uncondition=null_y,
-                            cfg_scale=guidance_scale,
-                            model_kwargs=model_kwargs,
-                        )
-                        sample = flow_solver.sample(
-                            z,
-                            steps=num_inference_steps,
-                        )
-                    elif self.vis_sampler == "flow_dpm-solver":
-                        scheduler = DPMS(
-                            self.model,
-                            condition=caption_embs,
-                            uncondition=null_y,
-                            guidance_type=self.guidance_type,
-                            cfg_scale=guidance_scale,
-                            pag_scale=pag_guidance_scale,
-                            pag_applied_layers=self.config.model.pag_applied_layers,
-                            model_type="flow",
-                            model_kwargs=model_kwargs,
-                            schedule="FLOW",
-                        )
-                        scheduler.register_progress_bar(self.progress_fn)
-                        sample = scheduler.sample(
-                            z,
-                            steps=num_inference_steps,
-                            order=2,
-                            skip_type="time_uniform_flow",
-                            method="multistep",
-                            flow_shift=self.flow_shift,
-                        )
+            with torch.no_grad():
+                # prepare text feature
+                if not self.config.text_encoder.chi_prompt:
+                    max_length_all = self.config.text_encoder.model_max_length
+                    prompts_all = prompts
+                else:
+                    chi_prompt = "\n".join(self.config.text_encoder.chi_prompt)
+                    prompts_all = [chi_prompt + prompt for prompt in prompts]
+                    num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt))
+                    max_length_all = (
+                        num_chi_prompt_tokens + self.config.text_encoder.model_max_length - 2
+                    )  # magic number 2: [bos], [_]
+
+                caption_token = self.tokenizer(
+                    prompts_all,
+                    max_length=max_length_all,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="pt",
+                ).to(device=self.device)
+                select_index = [0] + list(range(-self.config.text_encoder.model_max_length + 1, 0))
+                caption_embs = self.text_encoder(caption_token.input_ids, caption_token.attention_mask)[0][:, None][
+                    :, :, select_index
+                ].to(self.weight_dtype)
+                emb_masks = caption_token.attention_mask[:, select_index]
+                null_y = self.null_caption_embs.repeat(len(prompts), 1, 1)[:, None].to(self.weight_dtype)
+
+                n = len(prompts)
+                if latents is None:
+                    z = torch.randn(
+                        n,
+                        self.config.vae.vae_latent_dim,
+                        self.latent_size_h,
+                        self.latent_size_w,
+                        generator=generator,
+                        device=self.device,
+                        dtype=self.weight_dtype,
+                    )
+                else:
+                    z = latents.to(self.weight_dtype).to(self.device)
+                model_kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": ar}, mask=emb_masks)
+                if self.vis_sampler == "flow_euler":
+                    flow_solver = FlowEuler(
+                        self.model,
+                        condition=caption_embs,
+                        uncondition=null_y,
+                        cfg_scale=guidance_scale,
+                        model_kwargs=model_kwargs,
+                    )
+                    sample = flow_solver.sample(
+                        z,
+                        steps=num_inference_steps,
+                    )
+                elif self.vis_sampler == "flow_dpm-solver":
+                    scheduler = DPMS(
+                        self.model,
+                        condition=caption_embs,
+                        uncondition=null_y,
+                        guidance_type=self.guidance_type,
+                        cfg_scale=guidance_scale,
+                        pag_scale=pag_guidance_scale,
+                        pag_applied_layers=self.config.model.pag_applied_layers,
+                        model_type="flow",
+                        model_kwargs=model_kwargs,
+                        schedule="FLOW",
+                    )
+                    scheduler.register_progress_bar(self.progress_fn)
+                    sample = scheduler.sample(
+                        z,
+                        steps=num_inference_steps,
+                        order=2,
+                        skip_type="time_uniform_flow",
+                        method="multistep",
+                        flow_shift=self.flow_shift,
+                    )
 
             sample = sample.to(self.weight_dtype)
             with torch.no_grad():