Update app.py

app.py

@@ -1,12 +1,8 @@
-"""
-Gradio 应用：mini-DDPM
-
-"""
+# Gradio 应用：mini-DDPM
 
 import math
-import os
 import random
-from typing import Optional
+from typing import Optional, Dict
 
 import gradio as gr
 import numpy as np
@@ -25,6 +21,7 @@ CHECKPOINT_URL = (
 # 工具函数
 # -----------------------------
 
+
 def seed_all(seed: Optional[int] = 42):
     if seed is None:
         return
@@ -33,13 +30,12 @@ def seed_all(seed: Optional[int] = 42):
     torch.manual_seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
-    # 为了速度；如需完全复现，可改为 False。
+    # 为了速度；如需完全复现，可改为 False
     torch.backends.cudnn.benchmark = True
 
 
-
-_MODEL_CACHE: dict[str, torch.nn.Module] = {}
-_SCHED_CACHE: dict[str, DDPMSchedule] = {}
+_MODEL_CACHE: Dict[str, torch.nn.Module] = {}
+_SCHED_CACHE: Dict[str, DDPMSchedule] = {}
 
 
 def get_device() -> torch.device:
@@ -50,7 +46,7 @@ def load_model(device: torch.device) -> torch.nn.Module:
     key = str(device)
     if key in _MODEL_CACHE:
         return _MODEL_CACHE[key]
-    
+
     state_dict = torch.hub.load_state_dict_from_url(
         CHECKPOINT_URL, map_location=device, file_name="checkpoints.pt", progress=True
     )
@@ -74,13 +70,14 @@ def load_scheduler(device: torch.device) -> DDPMSchedule:
 # 采样核心
 # -----------------------------
 
+
 def sample_ddpm(
     num_samples: int = 9,
-    seed: int | None = 42,
-    progress: gr.Progress | None = None,
+    seed: Optional[int] = 42,
+    progress: Optional[gr.Progress] = None,
 ) -> Image.Image:
-
-    side = int(math.sqrt(num_samples))
+    # 将样本数调整为最接近的平方数，便于拼图展示
+    side = int(math.sqrt(max(1, num_samples)))
     num_samples = max(1, side * side)
 
     device = get_device()
@@ -110,12 +107,14 @@ def sample_ddpm(
     img = Image.fromarray((grid * 255).astype(np.uint8))
     return img
 
+
 # -----------------------------
 # Gradio UI
 # -----------------------------
 
+
 def ui_generate(num_samples, seed):
-    # seed == -1 表示使用随机种子
+    # seed < 0 或 None 表示使用随机种子
     real_seed = None if seed is None or int(seed) < 0 else int(seed)
     return sample_ddpm(
         num_samples=int(num_samples),
@@ -126,7 +125,6 @@ def ui_generate(num_samples, seed):
 
 def build_demo() -> gr.Blocks:
     with gr.Blocks(theme=gr.themes.Soft(), title="mini-DDPM") as demo:
-
         with gr.Row():
             with gr.Column(scale=2):
                 out_img = gr.Image(type="pil", label="采样结果")
@@ -138,8 +136,9 @@ def build_demo() -> gr.Blocks:
                     value=9,
                     label="样本数量",
                 )
-                                seed = gr.Number(value=42, precision=0, label="随机种子")
-                                btn = gr.Button("开始生成", variant="primary")
+                
+                seed = gr.Number(value=42, precision=0, label="随机种子")
+                btn = gr.Button("开始生成", variant="primary")
 
         btn.click(
             fn=ui_generate,