Update gradio_app.py

gradio_app.py

@@ -198,6 +198,122 @@ def _import_inference_bits():
 # -----------------------------------------------------------------------------
 SD15_PATH, _, _ = _download_models()
 
+# -----------------------------------------------------------------------------
+# Global model loading (CPU) so GPU task only does inference
+# -----------------------------------------------------------------------------
+def _resolve_trained_model_dir() -> str:
+    tm_dir = os.path.abspath("trained_model") if os.path.isdir("trained_model") else None
+    if tm_dir is None and os.path.isdir("pretrain"):
+        tm_dir = os.path.abspath("pretrain")
+    if tm_dir is None:
+        raise RuntimeError("Missing trained model weights. Provide TRAINED_MODEL_REPO or include ./pretrain.")
+    return tm_dir
+
+
+# Lazy globals
+G_ARGS = None
+G_INFER_CONFIG = None
+G_TOKENIZER = None
+G_IMAGE_ENCODER = None
+G_VAE = None
+G_UNET2 = None
+G_CONTROLNET = None
+G_DENOISING_UNET = None
+G_CC_PROJ = None
+G_HAIR_ENCODER = None
+
+
+def _load_models_cpu_once():
+    global G_ARGS, G_INFER_CONFIG, G_TOKENIZER, G_IMAGE_ENCODER, G_VAE
+    global G_UNET2, G_CONTROLNET, G_DENOISING_UNET, G_CC_PROJ, G_HAIR_ENCODER
+
+    if all(x is not None for x in (
+        G_ARGS, G_INFER_CONFIG, G_TOKENIZER, G_IMAGE_ENCODER, G_VAE,
+        G_UNET2, G_CONTROLNET, G_DENOISING_UNET, G_CC_PROJ, G_HAIR_ENCODER
+    )):
+        return
+
+    class _Args:
+        pretrained_model_name_or_path = SD15_PATH or os.path.abspath("stable-diffusion-v1-5/stable-diffusion-v1-5")
+        model_path = _resolve_trained_model_dir()
+        image_encoder = "openai/clip-vit-large-patch14"
+        controlnet_model_name_or_path = None
+        revision = None
+        output_dir = "gradio_outputs"
+        seed = 42
+        num_validation_images = 1
+        validation_ids = []
+        validation_hairs = []
+        use_fp16 = False
+        align_before_infer = True
+        align_size = 1024
+
+    G_ARGS = _Args()
+
+    # Import heavy libs only here
+    from test_stablehairv2 import AutoTokenizer, CLIPVisionModelWithProjection, AutoencoderKL, UNet2DConditionModel
+    from test_stablehairv2 import UNet3DConditionModel, CCProjection, ControlNetModel
+    from omegaconf import OmegaConf
+
+    # Config
+    G_INFER_CONFIG = OmegaConf.load('./configs/inference/inference_v2.yaml')
+
+    # Tokenizer / encoders / vae (CPU)
+    G_TOKENIZER = AutoTokenizer.from_pretrained(G_ARGS.pretrained_model_name_or_path, subfolder="tokenizer",
+                                                revision=G_ARGS.revision)
+    G_IMAGE_ENCODER = CLIPVisionModelWithProjection.from_pretrained(G_ARGS.image_encoder, revision=G_ARGS.revision)
+    G_VAE = AutoencoderKL.from_pretrained(G_ARGS.pretrained_model_name_or_path, subfolder="vae",
+                                          revision=G_ARGS.revision)
+
+    # UNet2D with 8-channel conv_in (CPU)
+    G_UNET2 = UNet2DConditionModel.from_pretrained(
+        G_ARGS.pretrained_model_name_or_path, subfolder="unet", revision=G_ARGS.revision, torch_dtype=torch.float32
+    )
+    conv_in_8 = torch.nn.Conv2d(8, G_UNET2.conv_in.out_channels, kernel_size=G_UNET2.conv_in.kernel_size,
+                                padding=G_UNET2.conv_in.padding)
+    conv_in_8.requires_grad_(False)
+    G_UNET2.conv_in.requires_grad_(False)
+    torch.nn.init.zeros_(conv_in_8.weight)
+    conv_in_8.weight[:, :4, :, :].copy_(G_UNET2.conv_in.weight)
+    conv_in_8.bias.copy_(G_UNET2.conv_in.bias)
+    G_UNET2.conv_in = conv_in_8
+
+    # ControlNet (CPU)
+    G_CONTROLNET = ControlNetModel.from_unet(G_UNET2)
+    state_dict2 = torch.load(os.path.join(G_ARGS.model_path, "pytorch_model.bin"), map_location="cpu")
+    G_CONTROLNET.load_state_dict(state_dict2, strict=False)
+
+    # UNet3D (CPU)
+    prefix = "motion_module"
+    ckpt_num = "4140000"
+    save_path = os.path.join(G_ARGS.model_path, f"{prefix}-{ckpt_num}.pth")
+    G_DENOISING_UNET = UNet3DConditionModel.from_pretrained_2d(
+        G_ARGS.pretrained_model_name_or_path,
+        save_path,
+        subfolder="unet",
+        unet_additional_kwargs=G_INFER_CONFIG.unet_additional_kwargs,
+    )
+
+    # CC projection (CPU)
+    G_CC_PROJ = CCProjection()
+    state_dict3 = torch.load(os.path.join(G_ARGS.model_path, "pytorch_model_1.bin"), map_location="cpu")
+    G_CC_PROJ.load_state_dict(state_dict3, strict=False)
+
+    # Hair encoder (CPU)
+    from ref_encoder.reference_unet import ref_unet
+    G_HAIR_ENCODER = ref_unet.from_pretrained(
+        G_ARGS.pretrained_model_name_or_path, subfolder="unet", revision=G_ARGS.revision, low_cpu_mem_usage=False,
+        device_map=None, ignore_mismatched_sizes=True
+    )
+    state_dict4 = torch.load(os.path.join(G_ARGS.model_path, "pytorch_model_2.bin"), map_location="cpu")
+    G_HAIR_ENCODER.load_state_dict(state_dict4, strict=False)
+
+
+try:
+    _load_models_cpu_once()
+except Exception as _e:
+    print(f"[init] Model preload warning: {_e}", flush=True)
+
 
 # -----------------------------------------------------------------------------
 # Gradio inference
@@ -211,7 +327,7 @@ def inference(id_image, hair_image):
     # ZeroGPU: 强制使用 'cuda' 设备（ZeroGPU 下 torch.cuda.is_available 可能为 False）。
     device = torch.device("cuda")
 
-    # 导入依赖
+    # 导入依赖（轻量函数，不再加载大模型）
     (
         log_validation,
         UNet3DConditionModel,
@@ -280,54 +396,15 @@ def inference(id_image, hair_image):
     )
     logger = logging.getLogger(__name__)
 
-    # Load tokenizer/encoders/vae
-    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer",
-                                              revision=args.revision)
-    image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder, revision=args.revision).to(device)
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae",
-                                        revision=args.revision).to(device, dtype=torch.float32)
-
-    from omegaconf import OmegaConf
-    infer_config = OmegaConf.load('./configs/inference/inference_v2.yaml')
-
-    # UNet2D with 8-channel conv_in
-    unet2 = UNet2DConditionModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, torch_dtype=torch.float32
-    ).to(device)
-    conv_in_8 = torch.nn.Conv2d(8, unet2.conv_in.out_channels, kernel_size=unet2.conv_in.kernel_size,
-                                padding=unet2.conv_in.padding)
-    conv_in_8.requires_grad_(False)
-    unet2.conv_in.requires_grad_(False)
-    torch.nn.init.zeros_(conv_in_8.weight)
-    conv_in_8.weight[:, :4, :, :].copy_(unet2.conv_in.weight)
-    conv_in_8.bias.copy_(unet2.conv_in.bias)
-    unet2.conv_in = conv_in_8
-
-    controlnet = ControlNetModel.from_unet(unet2).to(device)
-    state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model.bin"), map_location="cpu")
-    controlnet.load_state_dict(state_dict2, strict=False)
-
-    prefix = "motion_module"
-    ckpt_num = "4140000"
-    save_path = os.path.join(args.model_path, f"{prefix}-{ckpt_num}.pth")
-    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
-        args.pretrained_model_name_or_path,
-        save_path,
-        subfolder="unet",
-        unet_additional_kwargs=infer_config.unet_additional_kwargs,
-    ).to(device)
-
-    cc_projection = CCProjection().to(device)
-    state_dict3 = torch.load(os.path.join(args.model_path, "pytorch_model_1.bin"), map_location="cpu")
-    cc_projection.load_state_dict(state_dict3, strict=False)
-
-    from ref_encoder.reference_unet import ref_unet
-    Hair_Encoder = ref_unet.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, low_cpu_mem_usage=False,
-        device_map=None, ignore_mismatched_sizes=True
-    ).to(device)
-    state_dict4 = torch.load(os.path.join(args.model_path, "pytorch_model_2.bin"), map_location="cpu")
-    Hair_Encoder.load_state_dict(state_dict4, strict=False)
+    # 将已加载的全局模型迁移到 GPU
+    tokenizer = G_TOKENIZER
+    image_encoder = G_IMAGE_ENCODER.to(device)
+    vae = G_VAE.to(device, dtype=torch.float32)
+    unet2 = G_UNET2.to(device)
+    controlnet = G_CONTROLNET.to(device)
+    denoising_unet = G_DENOISING_UNET.to(device)
+    cc_projection = G_CC_PROJ.to(device)
+    Hair_Encoder = G_HAIR_ENCODER.to(device)
 
     # Run inference
     log_validation(