Add files using upload-large-folder tool

jit_diffusers/pipeline_jit.py

@@ -13,6 +13,21 @@ from .modeling_jit_transformer_2d import JiTTransformer2DModel
 from .scheduling_jit import JiTScheduler
 
 
+RECOMMENDED_CFG_BY_MODEL = {
+    "JiT-B/16": 3.0,
+    "JiT-L/16": 2.4,
+    "JiT-H/16": 2.2,
+    "JiT-B/32": 3.0,
+    "JiT-L/32": 2.5,
+    "JiT-H/32": 2.3,
+}
+
+RECOMMENDED_NOISE_BY_RESOLUTION = {
+    256: 1.0,
+    512: 2.0,
+}
+
+
 @dataclass
 class JiTPipelineOutput(BaseOutput):
     images: List["PIL.Image.Image"] | np.ndarray | torch.Tensor
@@ -51,10 +66,10 @@ class JiTPipeline(DiffusionPipeline):
         self,
         class_labels: int | List[int] | torch.Tensor,
         num_inference_steps: int = 50,
-        guidance_scale: float = 2.9,
+        guidance_scale: float | None = None,
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
-        noise_scale: float = 2.0,
+        noise_scale: float | None = None,
         t_eps: float = 5e-2,
         sampling_method: str | None = None,
         generator: torch.Generator | List[torch.Generator] | None = None,
@@ -81,6 +96,12 @@ class JiTPipeline(DiffusionPipeline):
         latent_size = int(self.transformer.config.sample_size)
         latent_channels = int(getattr(self.transformer.config, "in_channels", 3))
         num_classes = int(self.transformer.config.num_class_embeds)
+        model_type = str(getattr(self.transformer.config, "model_type", ""))
+
+        if guidance_scale is None:
+            guidance_scale = RECOMMENDED_CFG_BY_MODEL.get(model_type, 2.9)
+        if noise_scale is None:
+            noise_scale = RECOMMENDED_NOISE_BY_RESOLUTION.get(latent_size, 1.0)
 
         class_labels = class_labels.clamp(0, num_classes - 1)
         class_null = torch.full_like(class_labels, num_classes)
@@ -102,7 +123,9 @@ class JiTPipeline(DiffusionPipeline):
             x_uncond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_null).sample
             v_uncond = (x_uncond - z_value) / (1.0 - t).clamp_min(t_eps)
 
-            interval_mask = (t < guidance_interval_max) & (t > guidance_interval_min)
+            interval_mask = t < guidance_interval_max
+            if guidance_interval_min != 0.0:
+                interval_mask = interval_mask & (t > guidance_interval_min)
             scale = torch.where(
                 interval_mask,
                 torch.tensor(guidance_scale, device=self._execution_device, dtype=latents.dtype),

run_jit_diffusers_inference.py

@@ -11,6 +11,21 @@ if str(SCRIPT_DIR) not in sys.path:
 from jit_diffusers import JiTPipeline
 
 
+RECOMMENDED_CFG_BY_MODEL = {
+    "JiT-B/16": 3.0,
+    "JiT-L/16": 2.4,
+    "JiT-H/16": 2.2,
+    "JiT-B/32": 3.0,
+    "JiT-L/32": 2.5,
+    "JiT-H/32": 2.3,
+}
+
+RECOMMENDED_NOISE_BY_RESOLUTION = {
+    256: 1.0,
+    512: 2.0,
+}
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Run single-image JiT diffusers inference.")
     parser.add_argument("--model_path", type=str, required=True, help="Path to converted diffusers model directory.")
@@ -18,10 +33,20 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--class_label", type=int, default=207, help="ImageNet class id for conditional generation.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
     parser.add_argument("--steps", type=int, default=50, help="Number of ODE sampling steps.")
-    parser.add_argument("--cfg", type=float, default=2.9, help="Classifier-free guidance scale.")
+    parser.add_argument(
+        "--cfg",
+        type=float,
+        default=None,
+        help="Classifier-free guidance scale. Defaults to paper recommendation for the loaded model.",
+    )
     parser.add_argument("--interval_min", type=float, default=0.1, help="CFG interval min.")
     parser.add_argument("--interval_max", type=float, default=1.0, help="CFG interval max.")
-    parser.add_argument("--noise_scale", type=float, default=2.0, help="Initial Gaussian noise scale.")
+    parser.add_argument(
+        "--noise_scale",
+        type=float,
+        default=None,
+        help="Initial Gaussian noise scale. Defaults to paper recommendation for the loaded resolution.",
+    )
     parser.add_argument("--t_eps", type=float, default=5e-2, help="Small epsilon for timestep denominator.")
     parser.add_argument(
         "--device",
@@ -59,6 +84,14 @@ def resolve_dtype(name: str, device: torch.device) -> torch.dtype:
     return torch.float32
 
 
+def resolve_generation_defaults(pipe: JiTPipeline, cfg: float | None, noise_scale: float | None) -> tuple[float, float]:
+    model_type = str(getattr(pipe.transformer.config, "model_type", ""))
+    sample_size = int(getattr(pipe.transformer.config, "sample_size", 256))
+    resolved_cfg = cfg if cfg is not None else RECOMMENDED_CFG_BY_MODEL.get(model_type, 2.9)
+    resolved_noise_scale = noise_scale if noise_scale is not None else RECOMMENDED_NOISE_BY_RESOLUTION.get(sample_size, 1.0)
+    return resolved_cfg, resolved_noise_scale
+
+
 def main() -> None:
     args = parse_args()
     device = resolve_device(args.device)
@@ -70,15 +103,16 @@ def main() -> None:
     pipe.transformer = pipe.transformer.to(device=device, dtype=dtype)
     pipe.transformer.eval()
     sampling_method = None if args.solver == "scheduler" else args.solver
+    cfg, noise_scale = resolve_generation_defaults(pipe, args.cfg, args.noise_scale)
 
     generator = torch.Generator(device=device).manual_seed(args.seed)
     output = pipe(
         class_labels=[args.class_label],
         num_inference_steps=args.steps,
-        guidance_scale=args.cfg,
+        guidance_scale=cfg,
         guidance_interval_min=args.interval_min,
         guidance_interval_max=args.interval_max,
-        noise_scale=args.noise_scale,
+        noise_scale=noise_scale,
         t_eps=args.t_eps,
         sampling_method=sampling_method,
         generator=generator,
@@ -89,6 +123,7 @@ def main() -> None:
     output_path = Path(args.output_path)
     output_path.parent.mkdir(parents=True, exist_ok=True)
     image.save(output_path)
+    print(f"Used sampling hyperparameters: cfg={cfg}, noise_scale={noise_scale}")
     print(f"Saved image to: {output_path}")