scripts/inference_magicdrive.py

import os
import gc
import sys
import time
import copy
from pprint import pformat
from datetime import timedelta
from functools import partial

# Allow loading numpy 2.x pickles with numpy 1.x
import numpy as _np
sys.modules.setdefault("numpy._core", _np.core)

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
sys.path.append(".")
DEVICE_TYPE = os.environ.get("DEVICE_TYPE", "gpu")

import torch
if not torch.cuda.is_available() or DEVICE_TYPE == 'npu':
    USE_NPU = True
    os.environ['DEVICE_TYPE'] = "npu"
    DEVICE_TYPE = "npu"
    print("Enable NPU!")
    try:
        # just before torch_npu, let xformers know there is no gpu
        import xformers
        import xformers.ops
    except Exception as e:
        print(f"Got {e} during import xformers!")
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
else:
    USE_NPU = False
import magicdrivedit.utils.module_contrib

import colossalai
import torch.distributed as dist
from torch.utils.data import Subset
from einops import rearrange, repeat
from colossalai.cluster import DistCoordinator, ProcessGroupMesh
# NOTE: do NOT import from mmengine.runner — its import chain reaches
# torch.distributed.optim which triggers a PyTorch JIT bug on Blackwell
# (LaTeX in docstring → SyntaxError at ast.parse).
def set_random_seed(seed: int = 1024) -> None:
    import random
    random.seed(seed)
    import numpy as np
    np.random.seed(seed)
    import torch
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

from tqdm import tqdm
from hydra import compose, initialize
from omegaconf import OmegaConf
from mmcv.parallel import DataContainer

from magicdrivedit.acceleration.parallel_states import (
    set_sequence_parallel_group,
    get_sequence_parallel_group,
    set_data_parallel_group,
    get_data_parallel_group
)
from magicdrivedit.datasets import save_sample
from magicdrivedit.datasets.dataloader import prepare_dataloader
from magicdrivedit.datasets.dataloader import prepare_dataloader
from magicdrivedit.models.text_encoder.t5 import text_preprocessing
from magicdrivedit.registry import DATASETS, MODELS, SCHEDULERS, build_module
from magicdrivedit.utils.config_utils import parse_configs, define_experiment_workspace, save_training_config, merge_dataset_cfg, mmengine_conf_get, mmengine_conf_set
from magicdrivedit.utils.inference_utils import (
    apply_mask_strategy,
    get_save_path_name,
    concat_6_views_pt,
    add_null_condition,
    enable_offload,
)
from magicdrivedit.utils.misc import (
    reset_logger,
    is_distributed,
    is_main_process,
    to_torch_dtype,
    collate_bboxes_to_maxlen,
    move_to,
    add_box_latent,
)
from magicdrivedit.utils.train_utils import sp_vae


TILING_PARAM = {
    "default": dict(),  # it is designed for CogVideoX's 720x480, 4.5 GB
    "384": dict(  # about 14.2 GB
        tile_sample_min_height = 384,  # should be 48n
        tile_sample_min_width = 720,  # should be 40n
    ),
}


def set_omegaconf_key_value(cfg, key, value):
    p, m = key.rsplit(".", 1)
    node = cfg
    for pk in p.split("."):
        node = getattr(node, pk)
    node[m] = value

    
def main():
    torch.set_grad_enabled(False)
    # ======================================================
    # configs & runtime variables
    # ======================================================
    # == parse configs ==
    cfg = parse_configs(training=False)
    if cfg.get("vsdebug", False):
        import debugpy
        debugpy.listen(5678)
        print("Waiting for debugger attach")
        debugpy.wait_for_client()
        print('Attached, continue...')

    # == dataset config ==
    if cfg.num_frames is None:
        num_data_cfgs = len(cfg.data_cfg_names)
        datasets = []
        val_datasets = []
        for (res, data_cfg_name), overrides in zip(
                cfg.data_cfg_names, cfg.get("dataset_cfg_overrides", [[]] * num_data_cfgs)):
            dataset, val_dataset = merge_dataset_cfg(cfg, data_cfg_name, overrides)
            datasets.append((res, dataset))
            val_datasets.append((res, val_dataset))
        dataset = {"type": "NuScenesMultiResDataset", "cfg": datasets}
        val_dataset = {"type": "NuScenesMultiResDataset", "cfg": val_datasets}
    else:
        dataset, val_dataset = merge_dataset_cfg(
            cfg, cfg.data_cfg_name, cfg.get("dataset_cfg_overrides", []),
            cfg.num_frames)
    if cfg.get("use_train", False):
        cfg.dataset = dataset
        tag = cfg.get("tag", "")
        cfg.tag = "train" if tag == "" else f"{tag}_train"
    else:
        cfg.dataset = val_dataset
    # set img_collate_param
    if hasattr(cfg.dataset, "img_collate_param"):
        cfg.dataset.img_collate_param.is_train = False  # Important!
    else:
        for d in cfg.dataset.cfg:
            d[1].img_collate_param.is_train = False  # Important!
    cfg.batch_size = 1
    # for lower cpu memory in dataloading
    cfg.ignore_ori_imgs = cfg.get("ignore_ori_imgs", True)
    if cfg.ignore_ori_imgs:
        cfg.dataset.drop_ori_imgs = True

    # for lower gpu memory in vae decoding
    cfg.vae_tiling = cfg.get("vae_tiling", None)

    # edit annotations
    if cfg.get("allow_class", None) != None:
        cfg.dataset.allow_class = cfg.allow_class
    if cfg.get("del_box_ratio", None) != None:
        cfg.dataset.del_box_ratio = cfg.del_box_ratio
    if cfg.get("drop_nearest_car", None) != None:
        cfg.dataset.drop_nearest_car = cfg.drop_nearest_car

    # == device and dtype ==
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg_dtype = cfg.get("dtype", "bf16")
    assert cfg_dtype in ["fp16", "bf16", "fp32"], f"Unknown mixed precision {cfg_dtype}"
    dtype = to_torch_dtype(cfg.get("dtype", "bf16"))
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    if USE_NPU:  # disable some kernels
        if mmengine_conf_get(cfg, "text_encoder.shardformer", None):
            mmengine_conf_set(cfg, "text_encoder.shardformer", False)
        if mmengine_conf_get(cfg, "model.bbox_embedder_param.enable_xformers", None):
            mmengine_conf_set(cfg, "model.bbox_embedder_param.enable_xformers", False)
        if mmengine_conf_get(cfg, "model.frame_emb_param.enable_xformers", None):
            mmengine_conf_set(cfg, "model.frame_emb_param.enable_xformers", False)

    # == init distributed env ==
    if is_distributed():
        # colossalai.launch_from_torch({})
        dist.init_process_group(backend="nccl", timeout=timedelta(hours=1))
        torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
        cfg.sp_size = dist.get_world_size()
    else:
        dist.init_process_group(
            backend="nccl", world_size=1, rank=0,
            init_method="tcp://localhost:12355")
        cfg.sp_size = 1
    coordinator = DistCoordinator()
    if cfg.sp_size > 1:
        DP_AXIS, SP_AXIS = 0, 1
        dp_size = dist.get_world_size() // cfg.sp_size
        pg_mesh = ProcessGroupMesh(dp_size, cfg.sp_size)
        dp_group = pg_mesh.get_group_along_axis(DP_AXIS)
        sp_group = pg_mesh.get_group_along_axis(SP_AXIS)
        set_sequence_parallel_group(sp_group)
        print(f"Using sp_size={cfg.sp_size}")
    else:
        # TODO: sequence_parallel_group unset!
        dp_group = dist.group.WORLD
    set_data_parallel_group(dp_group)
    enable_sequence_parallelism = cfg.sp_size > 1
    set_random_seed(seed=cfg.get("seed", 1024))

    # == init exp_dir ==
    cfg.outputs = cfg.get("outputs", "outputs/test")
    exp_name, exp_dir = define_experiment_workspace(cfg, use_date=True)
    cfg.save_dir = os.path.join(exp_dir, "generation")
    coordinator.block_all()
    if coordinator.is_master():
        os.makedirs(exp_dir, exist_ok=True)
        save_training_config(cfg.to_dict(), exp_dir)
    coordinator.block_all()

    # == init logger ==
    logger = reset_logger(exp_dir)
    logger.info("Inference configuration:\n %s", pformat(cfg.to_dict()))
    verbose = cfg.get("verbose", 1)

    # ======================================================
    # 2. build dataset and dataloader
    # ======================================================
    if cfg.get("val", None):
        validation_index = cfg.val.validation_index
        if validation_index == "all":
            raise NotImplementedError()
        cfg.num_sample = cfg.val.get("num_sample", 1)
        cfg.scheduler = cfg.val.get("scheduler", cfg.scheduler)
    else:
        validation_index = cfg.get("validation_index", "all")

    # == build dataset ==
    logger.info("Building dataset...")
    dataset = build_module(cfg.dataset, DATASETS)
    if validation_index == "even":
        idxs = list(range(0, len(dataset), 2))
        dataset = torch.utils.data.Subset(dataset, idxs)
    elif validation_index == "odd":
        idxs = list(reversed(list(range(1, len(dataset), 2))))  # reversed!
        dataset = torch.utils.data.Subset(dataset, idxs)
    elif validation_index != "all":
        dataset = torch.utils.data.Subset(dataset, validation_index)
    logger.info(f"Your validation index: {validation_index}")
    logger.info("Dataset contains %s samples.", len(dataset))

    # == build dataloader ==
    dataloader_args = dict(
        dataset=dataset,
        batch_size=cfg.get("batch_size", 1),
        num_workers=cfg.get("num_workers", 0),
        seed=cfg.get("seed", 1024),
        shuffle=isinstance(validation_index, str),  # changed
        drop_last=False,  # changed
        pin_memory=True,
        process_group=get_data_parallel_group(),
        prefetch_factor=cfg.get("prefetch_factor", None),
    )
    dataloader, sampler = prepare_dataloader(
        bucket_config=cfg.get("bucket_config", None),
        num_bucket_build_workers=cfg.get("num_bucket_build_workers", 1),
        **dataloader_args,
    )
    num_steps_per_epoch = len(dataloader)

    def collate_data_container_fn(batch, *, collate_fn_map=None):
        return batch
    # add datacontainer handler
    torch.utils.data._utils.collate.default_collate_fn_map.update({
        DataContainer: collate_data_container_fn
    })

    # ======================================================
    # build model & load weights
    # ======================================================
    logger.info("Building models...")
    # == build text-encoder and vae ==
    # NOTE: set to true/false,
    # https://github.com/huggingface/transformers/issues/5486
    # if the program gets stuck, try set it to false
    os.environ['TOKENIZERS_PARALLELISM'] = "true"
    text_encoder = build_module(cfg.text_encoder, MODELS, device=device)
    vae = build_module(cfg.vae, MODELS).to(device, dtype).eval()
    if cfg.vae_tiling:
        vae.module.enable_tiling(**TILING_PARAM[str(cfg.vae_tiling)])
        logger.info(f"VAE Tiling is enabled with {TILING_PARAM[str(cfg.vae_tiling)]}")

    # == build diffusion model ==
    model = (
        build_module(
            cfg.model,
            MODELS,
            input_size=(None, None, None),
            in_channels=vae.out_channels,
            caption_channels=text_encoder.output_dim,
            model_max_length=text_encoder.model_max_length,
            enable_sequence_parallelism=enable_sequence_parallelism,
        )
        .to(device, dtype)
        .eval()
    )
    text_encoder.y_embedder = model.y_embedder  # HACK: for classifier-free guidance

    # == build scheduler ==
    scheduler = build_module(cfg.scheduler, SCHEDULERS)

    # ======================================================
    # inference
    # ======================================================
    cfg.cpu_offload = cfg.get("cpu_offload", False)
    if cfg.cpu_offload:
        text_encoder.t5.model.to("cpu")
        model.to("cpu")
        vae.to("cpu")
        text_encoder.t5.model, model, vae, last_hook = enable_offload(
            text_encoder.t5.model, model, vae, device)
    # == load prompts ==
    # prompts = cfg.get("prompt", None)
    start_idx = cfg.get("start_index", 0)

    # == prepare arguments ==
    batch_size = cfg.get("batch_size", 1)
    num_sample = cfg.get("num_sample", 1)

    save_dir = cfg.save_dir
    os.makedirs(save_dir, exist_ok=True)
    sample_name = cfg.get("sample_name", None)
    prompt_as_path = cfg.get("prompt_as_path", False)

    # == Iter over all samples ==
    start_step = 0
    assert batch_size == 1
    sampler.set_epoch(0)
    dataloader_iter = iter(dataloader)
    with tqdm(
        enumerate(dataloader_iter, start=start_step),
        desc=f"Generating",
        disable=not coordinator.is_master() or not verbose,
        initial=start_step,
        total=num_steps_per_epoch,
    ) as pbar:
        for i, batch in pbar:
            if cfg.ignore_ori_imgs:
                B, T, NC = 1, *batch["pixel_values_shape"][0].tolist()[:2]
                latent_size = vae.get_latent_size(
                    (T, *batch["pixel_values_shape"][0].tolist()[-2:]))
            else:
                B, T, NC = batch["pixel_values"].shape[:3]
                latent_size = vae.get_latent_size((T, *batch["pixel_values"].shape[-2:]))
                # == prepare batch prompts ==
                x = batch.pop("pixel_values").to(device, dtype)
                x = rearrange(x, "B T NC C ... -> (B NC) C T ...")  # BxNC, C, T, H, W
            y = batch.pop("captions")[0]  # B, just take first frame
            maps = batch.pop("bev_map_with_aux").to(device, dtype)  # B, T, C, H, W
            bbox = batch.pop("bboxes_3d_data")
            # B len list (T, NC, len, 8, 3)
            bbox = [bbox_i.data for bbox_i in bbox]
            # B, T, NC, len, 8, 3
            # TODO: `bbox` may have some redundancy on `NC` dim.
            # NOTE: we reshape the data later!
            bbox = collate_bboxes_to_maxlen(bbox, device, dtype, NC, T)
            # B, T, NC, 3, 7
            cams = batch.pop("camera_param").to(device, dtype)
            cams = rearrange(cams, "B T NC ... -> (B NC) T 1 ...")  # BxNC, T, 1, 3, 7
            rel_pos = batch.pop("frame_emb").to(device, dtype)
            rel_pos = repeat(rel_pos, "B T ... -> (B NC) T 1 ...", NC=NC)  # BxNC, T, 1, 4, 4

            # variable for inference
            batch_prompts = y
            # ms = mask_strategy[i : i + batch_size]
            ms = [""] * len(y)
            # refs = reference_path[i : i + batch_size]
            refs = [""] * len(y)

            # == model input format ==
            model_args = {}
            model_args["maps"] = maps
            model_args["bbox"] = bbox
            model_args["cams"] = cams
            model_args["rel_pos"] = rel_pos
            model_args["fps"] = batch.pop('fps')
            model_args['drop_cond_mask'] = torch.ones((B))  # camera
            model_args['drop_frame_mask'] = torch.ones((B, T))  # box & rel_pos
            model_args["height"] = batch.pop("height")
            model_args["width"] = batch.pop("width")
            model_args["num_frames"] = batch.pop("num_frames")
            model_args = move_to(model_args, device=device, dtype=dtype)
            # no need to move these
            model_args["mv_order_map"] = cfg.get("mv_order_map")
            model_args["t_order_map"] = cfg.get("t_order_map")

            # == Iter over number of sampling for one prompt ==
            save_fps = int(model_args['fps'][0])
            for ns in range(num_sample):
                gc.collect()
                torch.cuda.empty_cache()
                # == prepare save paths ==
                save_paths = [
                    get_save_path_name(
                        save_dir,
                        sample_name=sample_name,
                        sample_idx=start_idx + idx,
                        prompt=y[idx],
                        prompt_as_path=prompt_as_path,
                        num_sample=num_sample,
                        k=ns,
                    )
                    for idx in range(len(y))
                ]
                if cfg.get("force_daytime", False):
                    batch_prompts[0] = batch_prompts[0].lower()
                    batch_prompts[0] = "Daytime. " + batch_prompts[0]
                    # exclude rain
                    batch_prompts[0] = batch_prompts[0].replace("rain", "sunny")
                    batch_prompts[0] = batch_prompts[0].replace("water reflections", "")
                    batch_prompts[0] = batch_prompts[0].replace("reflections in water", "")
                    batch_prompts[0] = batch_prompts[0].replace(" with umbrellas", "")
                    batch_prompts[0] = batch_prompts[0].replace(" with umbrella", "")
                    batch_prompts[0] = batch_prompts[0].replace(" holds umbrella", "")
                    # exclude night
                    batch_prompts[0] = batch_prompts[0].replace("night", "")
                    batch_prompts[0] = batch_prompts[0].replace(" in dark", "")
                    batch_prompts[0] = batch_prompts[0].replace(" dark", "")
                    batch_prompts[0] = batch_prompts[0].replace(" difficult lighting", "")
                    # city
                    batch_prompts[0] = batch_prompts[0].replace("boston-seaport", "singapore-onenorth")
                    batch_prompts[0] = batch_prompts[0].replace("singapore-hollandvillage", "singapore-onenorth")
                    neg_prompts = ["Rain, Night, water reflections, umbrella"]
                elif cfg.get("force_rainy", False):
                    if "rain" not in batch_prompts[0].lower():
                        batch_prompts[0] = "A driving scene image at boston-seaport. Rain. water reflections."
                    neg_prompts = ["Daytime. night, onenorth, queenstown"]
                elif cfg.get("force_night", False):
                    if "night" not in batch_prompts[0].lower():
                        batch_prompts[0] = "A driving scene image at singapore-hollandvillage. Night, congestion. difficult lighting. very dark."
                    neg_prompts = ["Daytime. rain, boston-seaport"]
                else:
                    neg_prompts = None

                video_clips = []
                # == sampling ==
                torch.manual_seed(1024 + ns)  # NOTE: not sure how to handle loop, just change here.
                z = torch.randn(len(batch_prompts), vae.out_channels * NC, *latent_size, device=device, dtype=dtype)

                # == sample box ==
                if bbox is not None:
                    # null set values to all zeros, this should be safe
                    bbox = add_box_latent(bbox, B, NC, T, model.sample_box_latent)
                    # overwrite!
                    new_bbox = {}
                    for k, v in bbox.items():
                        new_bbox[k] = rearrange(v, "B T NC ... -> (B NC) T ...")  # BxNC, T, len, 3, 7
                    model_args["bbox"] = move_to(new_bbox, device=device, dtype=dtype)

                # == add null condition ==
                # y is handled by scheduler.sample
                if cfg.scheduler.type == "dpm-solver" and cfg.scheduler.cfg_scale == 1.0 or (
                    cfg.scheduler.type in ["rflow-slice",]
                ):
                    _model_args = copy.deepcopy(model_args)
                else:
                    _model_args = add_null_condition(
                        copy.deepcopy(model_args),
                        model.camera_embedder.uncond_cam.to(device),
                        model.frame_embedder.uncond_cam.to(device),
                        prepend=(cfg.scheduler.type == "dpm-solver"),
                    )

                # == inference ==
                masks = None
                masks = apply_mask_strategy(z, refs, ms, 0, align=None)
                samples = scheduler.sample(
                    model,
                    text_encoder,
                    z=z,
                    prompts=batch_prompts,
                    neg_prompts=neg_prompts,
                    device=device,
                    additional_args=_model_args,
                    progress=verbose >= 1,
                    mask=masks,
                )
                samples = rearrange(samples, "B (C NC) T ... -> (B NC) C T ...", NC=NC)
                if cfg.sp_size > 1:
                    samples = sp_vae(
                        samples.to(dtype),
                        partial(vae.decode, num_frames=_model_args["num_frames"]),
                        get_sequence_parallel_group(),
                    )
                else:
                    samples = vae.decode(samples.to(dtype), num_frames=_model_args["num_frames"])
                samples = rearrange(samples, "(B NC) C T ... -> B NC C T ...", NC=NC)
                if cfg.cpu_offload:
                    last_hook.offload()
                if is_main_process():
                    vid_samples = []
                    for sample in samples:
                        vid_samples.append(
                            concat_6_views_pt(sample, oneline=False)
                        )
                    samples = torch.stack(vid_samples, dim=0)
                    video_clips.append(samples)
                    del vid_samples
                del samples
                coordinator.block_all()

                # == save samples ==
                torch.cuda.empty_cache()
                if is_main_process():
                    for idx, batch_prompt in enumerate(batch_prompts):
                        if verbose >= 1:
                            logger.info(f"Prompt: {batch_prompt}")
                            if neg_prompts is not None:
                                logger.info(f"Neg-prompt: {neg_prompts[idx]}")
                        save_path = save_paths[idx]
                        video = [video_clips[0][idx]]
                        video = torch.cat(video, dim=1)
                        save_path = save_sample(
                            video,
                            fps=save_fps,
                            save_path=save_path,
                            high_quality=True,
                            verbose=verbose >= 2,
                            save_per_n_frame=cfg.get("save_per_n_frame", -1),
                            force_image=cfg.get("force_image", False),
                        )
                del video_clips
                coordinator.block_all()
            
            # save_gt
            if is_main_process() and not cfg.ignore_ori_imgs:
                torch.cuda.empty_cache()
                samples = rearrange(x, "(B NC) C T H W -> B NC C T H W", NC=NC)
                for idx, sample in enumerate(samples):
                    vid_sample = concat_6_views_pt(sample, oneline=False)
                    save_path = save_sample(
                        vid_sample,
                        fps=save_fps,
                        save_path=os.path.join(save_dir, f"gt_{start_idx + idx:04d}"),
                        high_quality=True,
                        verbose=verbose >= 2,
                        save_per_n_frame=cfg.get("save_per_n_frame", -1),
                        force_image=cfg.get("force_image", False),
                    )
                del samples, vid_sample
            coordinator.block_all()
            start_idx += len(batch_prompts)
    logger.info("Inference finished.")
    logger.info("Saved %s samples to %s", start_idx - cfg.get("start_index", 0), save_dir)
    coordinator.destroy()


# ============================================================
# In-process inference entry (for ZeroGPU / app.py, no subprocess)
# ============================================================
# Models are cached in _GLOBAL and reused across requests; only the
# dataset/dataloader and the sampling loop run per request. This avoids
# the ZeroGPU SIGSEGV caused by launching a subprocess inside @spaces.GPU.
#
# Differences from main():
#  - cfg loaded programmatically (read_config + merge_args), no argparse/CLI
#  - dist.init_process_group guarded (not idempotent) + env:// (no fixed port)
#  - models cached in _GLOBAL; second call skips build
#  - Hydra GlobalHydra cleared before merge_dataset_cfg (long-lived process)
#  - cfg.config ensured (define_experiment_workspace reads it)
#  - logger handlers not leaked across calls
#  - coordinator.destroy() NOT called per request
#  - returns save_sample()'s return value (the output mp4 path)
#
# main() is untouched, so the local torchrun path is unchanged.

_GLOBAL = {
    "text_encoder": None,
    "vae": None,
    "model": None,
    "scheduler": None,
    "init_done": False,
    "logger_inited": False,
    "last_hook": None,
}


def _clear_global_hydra():
    """Clear GlobalHydra so merge_dataset_cfg can be called repeatedly in a
    long-lived process. Safe no-op if never initialized."""
    try:
        from hydra.core.global_hydra import GlobalHydra
        if GlobalHydra.instance().is_initialized():
            GlobalHydra.instance().clear()
    except Exception:
        pass


def _init_once(cfg, device, dtype, enable_sequence_parallelism):
    """One-time init: dist + models + global registrations. Idempotent.
    Returns (text_encoder, vae, model, scheduler). Cached across requests."""
    if _GLOBAL["init_done"]:
        return (_GLOBAL["text_encoder"], _GLOBAL["vae"],
                _GLOBAL["model"], _GLOBAL["scheduler"])

    # == dist init (NOT idempotent — guard it) ==
    # main() uses fixed init_method="tcp://localhost:12355" which collides
    # across repeated/concurrent calls; use env:// with a unique port.
    if not dist.is_initialized():
        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
        os.environ.setdefault("MASTER_PORT", "29500")
        dist.init_process_group(
            backend="nccl", world_size=1, rank=0, init_method="env://")
    # prepare_dataloader calls get_data_parallel_group().size() — must be set.
    set_data_parallel_group(dist.group.WORLD)

    # == global registrations (idempotent, but only need once) ==
    os.environ['TOKENIZERS_PARALLELISM'] = "true"
    torch.utils.data._utils.collate.default_collate_fn_map.update({
        DataContainer: lambda batch, *, collate_fn_map=None: batch
    })

    # == build models ==
    text_encoder = build_module(cfg.text_encoder, MODELS, device=device)
    vae = build_module(cfg.vae, MODELS).to(device, dtype).eval()
    if cfg.get("vae_tiling", None):
        vae.module.enable_tiling(**TILING_PARAM[str(cfg.vae_tiling)])
    model = (
        build_module(
            cfg.model,
            MODELS,
            input_size=(None, None, None),
            in_channels=vae.out_channels,
            caption_channels=text_encoder.output_dim,
            model_max_length=text_encoder.model_max_length,
            enable_sequence_parallelism=enable_sequence_parallelism,
        )
        .to(device, dtype)
        .eval()
    )
    text_encoder.y_embedder = model.y_embedder  # HACK for CFG, idempotent
    scheduler = build_module(cfg.scheduler, SCHEDULERS)

    # cpu_offload hooks installed once (not per request) to avoid stacking.
    if cfg.get("cpu_offload", False):
        text_encoder.t5.model.to("cpu")
        model.to("cpu")
        vae.to("cpu")
        text_encoder.t5.model, model, vae, last_hook = enable_offload(
            text_encoder.t5.model, model, vae, device)
        _GLOBAL["last_hook"] = last_hook

    _GLOBAL.update(
        text_encoder=text_encoder, vae=vae, model=model,
        scheduler=scheduler, init_done=True,
    )
    return text_encoder, vae, model, scheduler


def _get_logger(exp_dir, verbose):
    """Logger that does not leak handlers across repeated calls.
    First call sets up stream + file handlers; later calls only swap the
    file handler to the new exp_dir."""
    import logging
    logger = logging.getLogger()
    level = logging.DEBUG if (verbose and verbose >= 2) else logging.INFO
    logger.setLevel(level)
    formatter = logging.Formatter(
        "[\033[34m%(asctime)s\033[0m][%(name)s][%(levelname)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    if not _GLOBAL["logger_inited"]:
        sh = logging.StreamHandler()
        sh.setFormatter(formatter)
        logger.addHandler(sh)
        _GLOBAL["logger_inited"] = True
    # refresh file handler to the new exp_dir each call
    for h in list(logger.handlers):
        if isinstance(h, logging.FileHandler):
            logger.removeHandler(h)
    fh = logging.FileHandler(f"{exp_dir}/log_0.txt")
    fh.setFormatter(formatter)
    logger.addHandler(fh)
    return logger


def run_inplace(
    config_path,
    num_frames,
    seed,
    pkl_path,
    *,
    dtype_str="fp16",
    cpu_offload=False,
    force_daytime=False,
    force_rainy=False,
    force_night=False,
    num_sampling_steps=None,
    progress_callback=None,
):
    """In-process single-scene inference. Returns the output mp4 path (str).

    Mirrors main()'s dataset build + sampling loop, but with cached models
    and the in-process guards above. Designed to be called directly inside a
    ZeroGPU @spaces.GPU function (no subprocess).
    """
    from argparse import Namespace
    from magicdrivedit.utils.config_utils import (
        read_config, merge_args, merge_dataset_cfg, define_experiment_workspace,
    )

    def _cb(p, d):
        if progress_callback is not None:
            try:
                progress_callback(p, d)
            except Exception:
                pass

    # ======================================================
    # B: load cfg programmatically (equivalent to --cfg-options)
    # ======================================================
    cfg = read_config(config_path)
    cfg_options = {
        "num_frames": num_frames,
        "seed": seed,
        "dtype": dtype_str,
        "cpu_offload": cpu_offload,
        "dataset_cfg_overrides": [
            ("dataset.data.val.ann_file", pkl_path),
        ],
    }
    if force_daytime:
        cfg_options["force_daytime"] = True
    if force_rainy:
        cfg_options["force_rainy"] = True
    if force_night:
        cfg_options["force_night"] = True
    if num_sampling_steps is not None:
        cfg_options["scheduler.num_sampling_steps"] = num_sampling_steps

    # Resolve model.from_pretrained: the inference configs ship with the
    # placeholder "???", which build_module treats as an HF repo id →
    # "Repo id must use alphanumeric chars..." error. Point it at the
    # locally-downloaded checkpoint.
    #
    # MUST be a .pt/.pth file (load_checkpoint path). The HF repo also has a
    # model/ dir of sharded pytorch_model-*.bin, but that is HF-format
    # sharding which load_checkpoint does NOT understand — it would look for
    # .../model/model and fail. So only accept a .pt file here.
    fp = cfg.model.get("from_pretrained", None) if hasattr(cfg, "model") else None
    if not fp or fp == "???":
        # Locate project root from this file's path (scripts/), not cwd.
        _root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        ckpt_root = os.path.join(_root, "ckpts", "MagicDriveDiT-stage3-40k-ft")
        ema_pt = os.path.join(ckpt_root, "ema.pt")
        if os.path.exists(ema_pt):
            cfg_options["model.from_pretrained"] = ema_pt
        else:
            # ema.pt missing — diagnose by listing what's actually there.
            import glob as _glob
            if os.path.isdir(ckpt_root):
                listing = _glob.glob(os.path.join(ckpt_root, "**"), recursive=True)
                listing_str = "\n  ".join(listing[:40]) or "(empty)"
            else:
                listing_str = "(ckpt dir does not exist)"
            raise FileNotFoundError(
                f"ema.pt not found at {ema_pt}. load_checkpoint needs a .pt "
                f"file (the model/ sharded dir is not directly loadable). "
                f"Re-run 'Download / Check Weights'. Contents:\n  {listing_str}")

    args = Namespace(ckpt_path=None, cfg_options=cfg_options)
    cfg = merge_args(cfg, args, training=False)

    # Disable apex FusedLayerNorm: the configs set enable_layernorm_kernel=True
    # (via `True and global_layernorm`), which requires apex.normalization.
    # FusedLayerNorm. apex is hard to build on the Space (source compile, CUDA
    # sensitive). Fall back to nn.LayerNorm — same math, marginally slower,
    # no result change. Applied recursively to all nested model cfgs.
    def _disable_layernorm_kernel(node):
        if isinstance(node, dict):
            for k, v in list(node.items()):
                if k == "enable_layernorm_kernel" and v:
                    node[k] = False
                else:
                    _disable_layernorm_kernel(v)
    _disable_layernorm_kernel(cfg)

    # define_experiment_workspace reads cfg.config (basename of config file).
    if not cfg.get("config", None):
        cfg.config = config_path

    # ======================================================
    # B: dataset config merge (Hydra) — clear GlobalHydra first
    # ======================================================
    _clear_global_hydra()
    if cfg.num_frames is None:
        num_data_cfgs = len(cfg.data_cfg_names)
        datasets = []
        val_datasets = []
        for (res, data_cfg_name), overrides in zip(
                cfg.data_cfg_names, cfg.get("dataset_cfg_overrides", [[]] * num_data_cfgs)):
            dataset, val_dataset = merge_dataset_cfg(cfg, data_cfg_name, overrides)
            datasets.append((res, dataset))
            val_datasets.append((res, val_dataset))
        dataset = {"type": "NuScenesMultiResDataset", "cfg": datasets}
        val_dataset = {"type": "NuScenesMultiResDataset", "cfg": val_datasets}
    else:
        dataset, val_dataset = merge_dataset_cfg(
            cfg, cfg.data_cfg_name, cfg.get("dataset_cfg_overrides", []),
            cfg.num_frames)
    if cfg.get("use_train", False):
        cfg.dataset = dataset
    else:
        cfg.dataset = val_dataset
    if hasattr(cfg.dataset, "img_collate_param"):
        cfg.dataset.img_collate_param.is_train = False  # Important!
    else:
        for d in cfg.dataset.cfg:
            d[1].img_collate_param.is_train = False  # Important!
    cfg.batch_size = 1
    cfg.ignore_ori_imgs = cfg.get("ignore_ori_imgs", True)
    if cfg.ignore_ori_imgs:
        cfg.dataset.drop_ori_imgs = True
    cfg.vae_tiling = cfg.get("vae_tiling", None)
    if cfg.get("allow_class", None) is not None:
        cfg.dataset.allow_class = cfg.allow_class
    if cfg.get("del_box_ratio", None) is not None:
        cfg.dataset.del_box_ratio = cfg.del_box_ratio
    if cfg.get("drop_nearest_car", None) is not None:
        cfg.dataset.drop_nearest_car = cfg.drop_nearest_car

    # ======================================================
    # device / dtype
    # ======================================================
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg_dtype = cfg.get("dtype", "bf16")
    assert cfg_dtype in ["fp16", "bf16", "fp32"], f"Unknown mixed precision {cfg_dtype}"
    dtype = to_torch_dtype(cfg_dtype)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    # ======================================================
    # dist + cached models
    # ======================================================
    enable_sequence_parallelism = False  # sp_size == 1
    _cb(0.02, "Initializing models...")
    text_encoder, vae, model, scheduler = _init_once(
        cfg, device, dtype, enable_sequence_parallelism)

    # ======================================================
    # exp_dir / save_dir (recomputed per request)
    # ======================================================
    cfg.outputs = cfg.get("outputs", "outputs/test")
    exp_name, exp_dir = define_experiment_workspace(cfg, use_date=True)
    cfg.save_dir = os.path.join(exp_dir, "generation")
    save_dir = cfg.save_dir
    os.makedirs(exp_dir, exist_ok=True)
    os.makedirs(save_dir, exist_ok=True)

    verbose = cfg.get("verbose", 1)
    logger = _get_logger(exp_dir, verbose)
    logger.info("Inference (in-process) configuration saved to %s", exp_dir)

    coordinator = DistCoordinator()

    # ======================================================
    # build dataset + dataloader (per request)
    # ======================================================
    _cb(0.05, "Loading scene data...")
    if cfg.get("val", None):
        validation_index = cfg.val.validation_index
        if validation_index == "all":
            raise NotImplementedError()
        cfg.num_sample = cfg.val.get("num_sample", 1)
        cfg.scheduler = cfg.val.get("scheduler", cfg.scheduler)
    else:
        validation_index = cfg.get("validation_index", "all")

    dataset = build_module(cfg.dataset, DATASETS)
    if validation_index == "even":
        idxs = list(range(0, len(dataset), 2))
        dataset = torch.utils.data.Subset(dataset, idxs)
    elif validation_index == "odd":
        idxs = list(reversed(list(range(1, len(dataset), 2))))
        dataset = torch.utils.data.Subset(dataset, idxs)
    elif validation_index != "all":
        dataset = torch.utils.data.Subset(dataset, validation_index)
    logger.info("Dataset contains %s samples.", len(dataset))

    dataloader_args = dict(
        dataset=dataset,
        batch_size=cfg.get("batch_size", 1),
        num_workers=cfg.get("num_workers", 0),
        seed=cfg.get("seed", 1024),
        shuffle=isinstance(validation_index, str),
        drop_last=False,
        pin_memory=True,
        process_group=get_data_parallel_group(),
        prefetch_factor=cfg.get("prefetch_factor", None),
    )
    dataloader, sampler = prepare_dataloader(
        bucket_config=cfg.get("bucket_config", None),
        num_bucket_build_workers=cfg.get("num_bucket_build_workers", 1),
        **dataloader_args,
    )
    num_steps_per_epoch = len(dataloader)

    # ======================================================
    # inference loop (mirrors main() L332-537)
    # ======================================================
    set_random_seed(seed=cfg.get("seed", 1024))
    start_idx = cfg.get("start_index", 0)
    batch_size = cfg.get("batch_size", 1)
    num_sample = cfg.get("num_sample", 1)
    sample_name = cfg.get("sample_name", None)
    prompt_as_path = cfg.get("prompt_as_path", False)

    start_step = 0
    assert batch_size == 1
    sampler.set_epoch(0)
    dataloader_iter = iter(dataloader)

    out_save_path = None
    _cb(0.10, "Generating...")
    with tqdm(
        enumerate(dataloader_iter, start=start_step),
        desc=f"Generating",
        disable=not coordinator.is_master() or not verbose,
        initial=start_step,
        total=num_steps_per_epoch,
    ) as pbar:
        for i, batch in pbar:
            if cfg.ignore_ori_imgs:
                B, T, NC = 1, *batch["pixel_values_shape"][0].tolist()[:2]
                latent_size = vae.get_latent_size(
                    (T, *batch["pixel_values_shape"][0].tolist()[-2:]))
            else:
                B, T, NC = batch["pixel_values"].shape[:3]
                latent_size = vae.get_latent_size((T, *batch["pixel_values"].shape[-2:]))
                x = batch.pop("pixel_values").to(device, dtype)
                x = rearrange(x, "B T NC C ... -> (B NC) C T ...")
            y = batch.pop("captions")[0]
            maps = batch.pop("bev_map_with_aux").to(device, dtype)
            bbox = batch.pop("bboxes_3d_data")
            bbox = [bbox_i.data for bbox_i in bbox]
            bbox = collate_bboxes_to_maxlen(bbox, device, dtype, NC, T)
            cams = batch.pop("camera_param").to(device, dtype)
            cams = rearrange(cams, "B T NC ... -> (B NC) T 1 ...")
            rel_pos = batch.pop("frame_emb").to(device, dtype)
            rel_pos = repeat(rel_pos, "B T ... -> (B NC) T 1 ...", NC=NC)

            batch_prompts = y
            ms = [""] * len(y)
            refs = [""] * len(y)

            model_args = {}
            model_args["maps"] = maps
            model_args["bbox"] = bbox
            model_args["cams"] = cams
            model_args["rel_pos"] = rel_pos
            model_args["fps"] = batch.pop('fps')
            model_args['drop_cond_mask'] = torch.ones((B))
            model_args['drop_frame_mask'] = torch.ones((B, T))
            model_args["height"] = batch.pop("height")
            model_args["width"] = batch.pop("width")
            model_args["num_frames"] = batch.pop("num_frames")
            model_args = move_to(model_args, device=device, dtype=dtype)
            model_args["mv_order_map"] = cfg.get("mv_order_map")
            model_args["t_order_map"] = cfg.get("t_order_map")

            save_fps = int(model_args['fps'][0])
            for ns in range(num_sample):
                gc.collect()
                torch.cuda.empty_cache()
                save_paths = [
                    get_save_path_name(
                        save_dir,
                        sample_name=sample_name,
                        sample_idx=start_idx + idx,
                        prompt=y[idx],
                        prompt_as_path=prompt_as_path,
                        num_sample=num_sample,
                        k=ns,
                    )
                    for idx in range(len(y))
                ]
                if cfg.get("force_daytime", False):
                    batch_prompts[0] = batch_prompts[0].lower()
                    batch_prompts[0] = "Daytime. " + batch_prompts[0]
                    batch_prompts[0] = batch_prompts[0].replace("rain", "sunny")
                    batch_prompts[0] = batch_prompts[0].replace("water reflections", "")
                    batch_prompts[0] = batch_prompts[0].replace("reflections in water", "")
                    batch_prompts[0] = batch_prompts[0].replace(" with umbrellas", "")
                    batch_prompts[0] = batch_prompts[0].replace(" with umbrella", "")
                    batch_prompts[0] = batch_prompts[0].replace(" holds umbrella", "")
                    batch_prompts[0] = batch_prompts[0].replace("night", "")
                    batch_prompts[0] = batch_prompts[0].replace(" in dark", "")
                    batch_prompts[0] = batch_prompts[0].replace(" dark", "")
                    batch_prompts[0] = batch_prompts[0].replace(" difficult lighting", "")
                    batch_prompts[0] = batch_prompts[0].replace("boston-seaport", "singapore-onenorth")
                    batch_prompts[0] = batch_prompts[0].replace("singapore-hollandvillage", "singapore-onenorth")
                    neg_prompts = ["Rain, Night, water reflections, umbrella"]
                elif cfg.get("force_rainy", False):
                    if "rain" not in batch_prompts[0].lower():
                        batch_prompts[0] = "A driving scene image at boston-seaport. Rain. water reflections."
                    neg_prompts = ["Daytime. night, onenorth, queenstown"]
                elif cfg.get("force_night", False):
                    if "night" not in batch_prompts[0].lower():
                        batch_prompts[0] = "A driving scene image at singapore-hollandvillage. Night, congestion. difficult lighting. very dark."
                    neg_prompts = ["Daytime. rain, boston-seaport"]
                else:
                    neg_prompts = None

                video_clips = []
                _cb(0.15, "Sampling...")
                torch.manual_seed(1024 + ns)
                z = torch.randn(len(batch_prompts), vae.out_channels * NC, *latent_size, device=device, dtype=dtype)

                if bbox is not None:
                    bbox = add_box_latent(bbox, B, NC, T, model.sample_box_latent)
                    new_bbox = {}
                    for k, v in bbox.items():
                        new_bbox[k] = rearrange(v, "B T NC ... -> (B NC) T ...")
                    model_args["bbox"] = move_to(new_bbox, device=device, dtype=dtype)

                if cfg.scheduler.type == "dpm-solver" and cfg.scheduler.cfg_scale == 1.0 or (
                    cfg.scheduler.type in ["rflow-slice",]
                ):
                    _model_args = copy.deepcopy(model_args)
                else:
                    _model_args = add_null_condition(
                        copy.deepcopy(model_args),
                        model.camera_embedder.uncond_cam.to(device),
                        model.frame_embedder.uncond_cam.to(device),
                        prepend=(cfg.scheduler.type == "dpm-solver"),
                    )

                masks = None
                masks = apply_mask_strategy(z, refs, ms, 0, align=None)
                samples = scheduler.sample(
                    model,
                    text_encoder,
                    z=z,
                    prompts=batch_prompts,
                    neg_prompts=neg_prompts,
                    device=device,
                    additional_args=_model_args,
                    progress=verbose >= 1,
                    mask=masks,
                )
                _cb(0.90, "Decoding VAE...")
                samples = rearrange(samples, "B (C NC) T ... -> (B NC) C T ...", NC=NC)
                # sp_size == 1 here, so plain vae.decode
                samples = vae.decode(samples.to(dtype), num_frames=_model_args["num_frames"])
                samples = rearrange(samples, "(B NC) C T ... -> B NC C T ...", NC=NC)
                if cfg.cpu_offload and _GLOBAL["last_hook"] is not None:
                    _GLOBAL["last_hook"].offload()
                if is_main_process():
                    vid_samples = []
                    for sample in samples:
                        vid_samples.append(concat_6_views_pt(sample, oneline=False))
                    samples = torch.stack(vid_samples, dim=0)
                    video_clips.append(samples)
                    del vid_samples
                del samples
                coordinator.block_all()

                torch.cuda.empty_cache()
                if is_main_process():
                    _cb(0.95, "Saving video...")
                    for idx, batch_prompt in enumerate(batch_prompts):
                        if verbose >= 1:
                            logger.info(f"Prompt: {batch_prompt}")
                            if neg_prompts is not None:
                                logger.info(f"Neg-prompt: {neg_prompts[idx]}")
                        save_path = save_paths[idx]
                        video = [video_clips[0][idx]]
                        video = torch.cat(video, dim=1)
                        save_path = save_sample(
                            video,
                            fps=save_fps,
                            save_path=save_path,
                            high_quality=True,
                            verbose=verbose >= 2,
                            save_per_n_frame=cfg.get("save_per_n_frame", -1),
                            force_image=cfg.get("force_image", False),
                        )
                        out_save_path = save_path  # capture return value
                del video_clips
                coordinator.block_all()

            # save_gt (only if not ignoring ori imgs)
            if is_main_process() and not cfg.ignore_ori_imgs:
                torch.cuda.empty_cache()
                samples = rearrange(x, "(B NC) C T H W -> B NC C T H W", NC=NC)
                for idx, sample in enumerate(samples):
                    vid_sample = concat_6_views_pt(sample, oneline=False)
                    save_sample(
                        vid_sample,
                        fps=save_fps,
                        save_path=os.path.join(save_dir, f"gt_{start_idx + idx:04d}"),
                        high_quality=True,
                        verbose=verbose >= 2,
                        save_per_n_frame=cfg.get("save_per_n_frame", -1),
                        force_image=cfg.get("force_image", False),
                    )
                del samples, vid_sample
            coordinator.block_all()
            start_idx += len(batch_prompts)

    logger.info("Inference (in-process) finished. Saved to %s", out_save_path)
    # NOTE: do NOT coordinator.destroy() — models are cached across requests.
    _cb(1.0, "Done!")
    return out_save_path


if __name__ == "__main__":
    main()