Function infer

examples/inference/stable_diffusion/sd3_generation.py:23–56

(args)

def infer(args):
    # ==============================
    # Launch colossalai, setup distributed environment
    # ==============================
    colossalai.launch_from_torch()
    coordinator = DistCoordinator()

    # ==============================
    # Load model and tokenizer
    # ==============================
    model_path_or_name = args.model
    model = MODEL_CLS.from_pretrained(model_path_or_name, torch_dtype=TORCH_DTYPE_MAP.get(args.dtype, None))

    # ==============================
    # Initialize InferenceEngine
    # ==============================
    coordinator.print_on_master(f"Initializing Inference Engine...")
    inference_config = InferenceConfig(
        dtype=args.dtype,
        max_batch_size=args.max_batch_size,
        tp_size=args.tp_size,
        use_cuda_kernel=args.use_cuda_kernel,
        patched_parallelism_size=dist.get_world_size(),
    )
    engine = InferenceEngine(model, inference_config=inference_config, verbose=True)

    # ==============================
    # Generation
    # ==============================
    coordinator.print_on_master(f"Generating...")
    out = engine.generate(prompts=[args.prompt], generation_config=DiffusionGenerationConfig())[0]
    if dist.get_rank() == 0:
        out.save(f"cat_parallel_size{dist.get_world_size()}.jpg")
    coordinator.print_on_master(out)


# colossalai run --nproc_per_node 1 examples/inference/stable_diffusion/sd3_generation.py -m MODEL_PATH