Delete demo_inference.py

demo_inference.py

@@ -1,161 +0,0 @@
-#!/usr/bin/env python3
-"""Smoke-test MVSplit-DiT inference from the converted Diffusers Hub folder."""
-
-from __future__ import annotations
-
-import argparse
-import importlib.util
-import sys
-from pathlib import Path
-
-import torch
-from diffusers import AutoencoderKLFlux2
-from transformers import AutoModel, AutoTokenizer
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run MVSplit-DiT inference.")
-    parser.add_argument(
-        "--model",
-        type=Path,
-        default=Path(__file__).resolve().parent,
-        help="Path to MVSplit-DiT-1000L pipeline directory.",
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        default="a red panda climbing a bamboo stalk",
-        help="Text prompt for generation.",
-    )
-    parser.add_argument("--height", type=int, default=256)
-    parser.add_argument("--width", type=int, default=256)
-    parser.add_argument("--num-inference-steps", type=int, default=35)
-    parser.add_argument("--guidance-scale", type=float, default=2.0)
-    parser.add_argument("--time-shift-alpha", type=float, default=4.0)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument(
-        "--output",
-        type=Path,
-        default=Path(__file__).resolve().parent / "demo.png",
-        help="Output image path. Ignored when --output-type=latent.",
-    )
-    parser.add_argument(
-        "--output-type",
-        choices=("pil", "latent"),
-        default="pil",
-        help="Return decoded image or raw latents.",
-    )
-    parser.add_argument(
-        "--skip-vae",
-        action="store_true",
-        help="Skip VAE decode even when output-type=pil (saves memory).",
-    )
-    parser.add_argument(
-        "--device",
-        choices=("auto", "cuda", "cpu"),
-        default="auto",
-        help="Execution device. auto prefers CUDA when available.",
-    )
-    parser.add_argument(
-        "--cpu-offload",
-        action="store_true",
-        help="Use sequential CPU offload instead of keeping the pipeline on GPU.",
-    )
-    return parser.parse_args()
-
-
-def _resolve_device(choice: str) -> torch.device:
-    if choice == "auto":
-        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    return torch.device(choice)
-
-
-def _load_pipeline_class(model_dir: Path):
-    transformer_path = model_dir / "transformer" / "transformer_mvsplit_dit.py"
-    spec = importlib.util.spec_from_file_location("transformer_mvsplit_dit", transformer_path)
-    module = importlib.util.module_from_spec(spec)
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-
-    pipe_spec = importlib.util.spec_from_file_location("mvsplit_pipeline", model_dir / "pipeline.py")
-    pipe_module = importlib.util.module_from_spec(pipe_spec)
-    sys.modules[pipe_spec.name] = pipe_module
-    pipe_spec.loader.exec_module(pipe_module)
-    return module.MVSplitDiTTransformer2DModel, pipe_module.MVSplitDiTPipeline
-
-
-def main() -> None:
-    args = parse_args()
-    model_dir = args.model.resolve()
-    device = _resolve_device(args.device)
-    transformer_cls, pipeline_cls = _load_pipeline_class(model_dir)
-
-    print(f"Loading components on {device}...", flush=True)
-    transformer = transformer_cls.from_pretrained(
-        model_dir / "transformer",
-        torch_dtype=torch.bfloat16,
-        local_files_only=True,
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer", local_files_only=True)
-    text_encoder = AutoModel.from_pretrained(
-        model_dir / "text_encoder",
-        torch_dtype=torch.bfloat16,
-        local_files_only=True,
-    )
-
-    vae = None
-    if not args.skip_vae and args.output_type == "pil":
-        vae = AutoencoderKLFlux2.from_pretrained(
-            model_dir / "vae",
-            torch_dtype=torch.bfloat16,
-            local_files_only=True,
-        )
-
-    pipe = pipeline_cls(
-        transformer=transformer,
-        scheduler=None,
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        time_shift_alpha=args.time_shift_alpha,
-    )
-    if args.cpu_offload and device.type == "cuda":
-        pipe.enable_sequential_cpu_offload(gpu_id=device.index or 0)
-    else:
-        pipe.to(device)
-
-    print(
-        f"Running inference ({args.num_inference_steps} steps, {args.height}x{args.width})...",
-        flush=True,
-    )
-    generator_device = "cpu" if args.cpu_offload else device.type
-    generator = torch.Generator(device=generator_device).manual_seed(args.seed)
-    result = pipe(
-        prompt=args.prompt,
-        height=args.height,
-        width=args.width,
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        generator=generator,
-        output_type=args.output_type,
-    )
-
-    if args.output_type == "latent":
-        latents = result.images
-        print(f"latent shape={tuple(latents.shape)} dtype={latents.dtype}")
-        print(
-            "latent stats:",
-            f"min={float(latents.min()):.4f}",
-            f"max={float(latents.max()):.4f}",
-            f"mean={float(latents.mean()):.4f}",
-        )
-        return
-
-    image = result.images[0]
-    args.output.parent.mkdir(parents=True, exist_ok=True)
-    image.save(args.output)
-    print(f"Saved image to {args.output}")
-
-
-if __name__ == "__main__":
-    main()