use pre-quantized INT8 model (no FP8 casting, no LoRA merge at runtime)

app.py

@@ -1,4 +1,4 @@
-"""FE2E: Depth + Normal estimation from a single image (CPU)"""
+"""FE2E: Depth + Normal estimation from a single image (CPU, pre-quantized INT8)"""
 from __future__ import annotations
 
 import gc
@@ -35,64 +35,42 @@ def _download(repo, filename, token=None):
 
 
 def _load_generator():
-    """Load FP8 model, cast to FP32, merge LoRA, quantize INT8."""
-    from safetensors.torch import load_file
+    """Load pre-quantized INT8 DiT + FP32 VAE. No FP8 casting, no LoRA merge at runtime."""
     from modules.model_edit import Step1XParams, Step1XEdit
     from modules.autoencoder import AutoEncoder
-    from infer.inference import ImageGenerator, load_state_dict, equip_dit_with_lora_sd_scripts
-    import numpy as np
+    from infer.inference import ImageGenerator
 
     token = os.environ.get("HF_TOKEN")
-    dit_path = _download("rkfg/Step1X-Edit-FP8", "step1x-edit-i1258-FP8.safetensors", token)
-    vae_path = _download("exander/FE2E", "pretrain/vae.safetensors", token)
-    lora_path = _download("exander/FE2E", "LDRN.safetensors", token)
+    dit_path = _download("WeReCooking2/FE2E-INT8", "dit_int8.pt", token)
+    vae_path = _download("WeReCooking2/FE2E-INT8", "vae_fp32.pt", token)
 
     args = Args()
-
-    print("[init] Building model on meta device...")
-    with torch.device("meta"):
-        ae = AutoEncoder(
-            resolution=256, in_channels=3, ch=128, out_ch=3,
-            ch_mult=[1, 2, 4, 4], num_res_blocks=2, z_channels=16,
-            scale_factor=0.3611, shift_factor=0.1159,
-        )
-        step1x_params = Step1XParams(
-            in_channels=64, out_channels=64, vec_in_dim=768, context_in_dim=4096,
-            hidden_size=3072, mlp_ratio=4.0, num_heads=24, depth=19,
-            depth_single_blocks=38, axes_dim=[16, 56, 56], theta=10_000, qkv_bias=True,
-        )
-        dit = Step1XEdit(step1x_params)
-
-    # Load weights as FP8, then cast to FP32 layer by layer to avoid 46 GB peak
-    print("[init] Loading FP8 weights and casting to FP32 (layer by layer)...")
+    print("[init] Loading pre-quantized INT8 DiT...")
     t0 = time.time()
-    fp8_sd = load_file(dit_path, device="cpu")
-    dit_sd = {}
-    for k, v in fp8_sd.items():
-        dit_sd[k] = v.float()
-    del fp8_sd
-    gc.collect()
 
-    dit = dit.to(dtype=torch.float32)
-    missing, unexpected = dit.load_state_dict(dit_sd, strict=False, assign=True)
-    del dit_sd
-    gc.collect()
-    print(f"[init] DiT loaded in {time.time()-t0:.0f}s (missing={len(missing)}, unexpected={len(unexpected)})")
+    # Build model structure, then quantize empty model, then load pre-quantized weights
+    step1x_params = Step1XParams(
+        in_channels=64, out_channels=64, vec_in_dim=768, context_in_dim=4096,
+        hidden_size=3072, mlp_ratio=4.0, num_heads=24, depth=19,
+        depth_single_blocks=38, axes_dim=[16, 56, 56], theta=10_000, qkv_bias=True,
+    )
+    dit = Step1XEdit(step1x_params).float()
+    dit = torch.quantization.quantize_dynamic(dit, {torch.nn.Linear}, dtype=torch.qint8)
+    dit.load_state_dict(torch.load(dit_path, map_location="cpu", weights_only=False))
+    print(f"[init] DiT loaded in {time.time()-t0:.0f}s")
 
     # Load VAE
-    ae = load_state_dict(ae, vae_path, "cpu")
-    ae = ae.float()
+    print("[init] Loading VAE (FP32)...")
+    ae = AutoEncoder(
+        resolution=256, in_channels=3, ch=128, out_ch=3,
+        ch_mult=[1, 2, 4, 4], num_res_blocks=2, z_channels=16,
+        scale_factor=0.3611, shift_factor=0.1159,
+    ).float()
+    ae.load_state_dict(torch.load(vae_path, map_location="cpu", weights_only=False))
 
-    # Merge LoRA in FP32 (full precision merge)
-    print("[init] Merging LoRA...")
-    equip_dit_with_lora_sd_scripts(ae, [None], dit, lora_path, device="cpu")
-
-    # Dynamic INT8 quantization
-    print("[init] Applying dynamic INT8 quantization...")
-    dit = torch.quantization.quantize_dynamic(dit, {torch.nn.Linear}, dtype=torch.qint8)
     gc.collect()
 
-    # Build generator wrapper
+    # Build generator wrapper (no LoRA needed, already merged)
     generator = ImageGenerator.__new__(ImageGenerator)
     generator.device = torch.device("cpu")
     generator.args = args
@@ -103,7 +81,7 @@ def _load_generator():
     generator.offload = False
     generator.lora_module = None
 
-    print("[init] Model ready for inference")
+    print(f"[init] Ready. Total load: {time.time()-t0:.0f}s")
     return generator
 
 
@@ -139,7 +117,7 @@ def generate(image):
     print(f"[gen] Input: {image.size}")
     t0 = time.time()
 
-    with torch.inference_mode(), torch.no_grad():
+    with torch.inference_mode():
         images, Lpred, Rpred = GENERATOR.generate_image(
             prompt="",
             negative_prompt="",
@@ -186,7 +164,7 @@ def main():
     with gr.Blocks(title="FE2E: Depth + Normal (CPU)") as demo:
         gr.Markdown(
             "**[FE2E](https://github.com/AMAP-ML/FE2E)** Depth + Normal from a single image (CVPR 2026). "
-            "Single denoise step, Step1X-Edit DiT + LDRN LoRA, dynamic INT8 on CPU."
+            "Single denoise step, Step1X-Edit DiT + LDRN LoRA (pre-merged), INT8 quantized on CPU."
         )
         with gr.Row():
             with gr.Column():