跑数据集遇到的问题

[liyanboSustech]

我发现拿taylorseer方法跑一整个coco数据集，前面生成的图片还可以接受，但后面生成的图片质量会非常差，不知道有没有人也遇到过这个问题

[Shenyi-Z]

Your description is too vague. What exactly is the "front," and what is the "back"? Could you provide the script and code that you ran?

[liyanboSustech]

These two script are both based on your application of TaylorSeer in FLUX and has been modified to support running multiple prompts.

import argparse
import logging
import os
import sys
from pathlib import Path
import time
import torch
import torch.distributed as dist
from transformers import T5EncoderModel

from xfuser import xFuserFluxPipeline, xFuserArgs
from xfuser.config import FlexibleArgumentParser
from xfuser.core.distributed import (
    get_world_group,
    get_data_parallel_rank,
    get_data_parallel_world_size,
    get_runtime_state,
    is_dp_last_group,
)

from xfuser.model_executor.models.transformers.transformer_flux import (
    xFuserFluxTransformer2DWrapper,  # noqa: F401
)

from typing import Any, Dict, List, Optional, Tuple, Union
from diffusers import DiffusionPipeline
from diffusers.models import FluxTransformer2DModel
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
import torch
import numpy as np
from xfuser.core.distributed.parallel_state import (
    get_tensor_model_parallel_world_size,
    is_pipeline_first_stage,
    is_pipeline_last_stage,
)

from forwards import taylorseer_flux_single_block_forward, taylorseer_flux_double_block_forward, taylorseer_flux_forward, taylorseer_xfuser_flux_forward


def _normalize_prompts(prompts):
    if isinstance(prompts, list):
        return prompts
    if isinstance(prompts, tuple):
        return list(prompts)
    if prompts is None:
        return []
    return [prompts]


def _chunk_prompts(prompts: List[str], batch_size: int) -> List[List[str]]:
    if batch_size <= 0:
        batch_size = 1
    if not prompts:
        return []
    return [prompts[i : i + batch_size] for i in range(0, len(prompts), batch_size)]


def _load_prompt_batches(
    prompt_file: Path, batch_size: int, limit: Optional[int] = None
) -> List[List[str]]:
    if not prompt_file.exists():
        raise FileNotFoundError(f"Prompt file not found: {prompt_file}")
    with prompt_file.open("r", encoding="utf-8") as handle:
        prompts = [line.strip() for line in handle if line.strip()]
    if limit is not None and limit > 0:
        prompts = prompts[:limit]
    return _chunk_prompts(prompts, batch_size)

def main():
    prompt_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
    prompt_parser.add_argument(
        "--prompt-file",
        dest="prompt_file",
        type=str,
        default=None,
        help="Optional text file providing prompts, one per line.",
    )
    prompt_parser.add_argument(
        "--prompt-file-batch-size",
        dest="prompt_file_batch_size",
        type=int,
        default=None,
        help="Batch size to use when consuming prompts from --prompt-file.",
    )
    prompt_parser.add_argument(
        "--output-dir",
        dest="output_dir",
        type=str,
        default="./results",
        help="Directory to store generated images.",
    )
    prompt_args, remaining_argv = prompt_parser.parse_known_args()

    parser = FlexibleArgumentParser(description="xFuser Arguments")
    args = xFuserArgs.add_cli_args(parser).parse_args(remaining_argv)
    engine_args = xFuserArgs.from_cli_args(args)
    prompt_file_path = (
        Path(prompt_args.prompt_file).expanduser()
        if prompt_args.prompt_file
        else None
    )
    prompt_file_batch_size = prompt_args.prompt_file_batch_size
    output_dir = Path(prompt_args.output_dir).expanduser()
    engine_config, input_config = engine_args.create_config()
    engine_config.runtime_config.dtype = torch.bfloat16
    local_rank = get_world_group().local_rank
    text_encoder_2 = T5EncoderModel.from_pretrained(engine_config.model_config.model, subfolder="text_encoder_2", torch_dtype=torch.bfloat16)

    if args.use_fp8_t5_encoder:
        from optimum.quanto import freeze, qfloat8, quantize
        logging.info(f"rank {local_rank} quantizing text encoder 2")
        quantize(text_encoder_2, weights=qfloat8)
        freeze(text_encoder_2)

    pipe = xFuserFluxPipeline.from_pretrained(
        pretrained_model_name_or_path=engine_config.model_config.model,
        engine_config=engine_config,
        torch_dtype=torch.bfloat16,
        text_encoder_2=text_encoder_2,
    )

    if args.enable_sequential_cpu_offload:
        pipe.enable_sequential_cpu_offload(gpu_id=local_rank)
        logging.info(f"rank {local_rank} sequential CPU offload enabled")
    else:
        pipe = pipe.to(f"cuda:{local_rank}")

    if not dist.is_initialized() or get_tensor_model_parallel_world_size() == 1:
        if not isinstance(pipe.transformer, xFuserFluxTransformer2DWrapper):
            # 
            original_transformer = pipe.transformer
            
            pipe.transformer = xFuserFluxTransformer2DWrapper(original_transformer)

            
    pipe.transformer.__class__.num_steps = input_config.num_inference_steps
    pipe.transformer.__class__.forward = taylorseer_xfuser_flux_forward

    for double_transformer_block in pipe.transformer.transformer_blocks:
        double_transformer_block.__class__.forward = taylorseer_flux_double_block_forward

    for single_transformer_block in pipe.transformer.single_transformer_blocks:
        single_transformer_block.__class__.forward = taylorseer_flux_single_block_forward
    
    joint_attention_kwargs = {}

    prompt_batches: Optional[List[List[str]]] = None
    total_prompts = 0
    if prompt_file_path:
        batch_size_from_file = prompt_file_batch_size or input_config.batch_size or 1
        if batch_size_from_file <= 0:
            batch_size_from_file = 1
        prompt_batches = _load_prompt_batches(prompt_file_path, batch_size_from_file)
        total_prompts = sum(len(batch) for batch in prompt_batches)
        if not prompt_batches:
            logging.warning(
                "No prompts found in %s, falling back to CLI prompt.", prompt_file_path
            )

    parameter_peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}")

    output_dir.mkdir(parents=True, exist_ok=True)

    prepared_batch_size = input_config.batch_size
    pipe.prepare_run(input_config, steps=1)

    def ensure_prepared(batch_size: int):
        nonlocal prepared_batch_size
        if batch_size <= 0:
            batch_size = 1
        if batch_size != prepared_batch_size:
            input_config.batch_size = batch_size
            pipe.prepare_run(input_config, steps=1)
            prepared_batch_size = batch_size

    torch.cuda.reset_peak_memory_stats()

    parallel_info = (
        f"dp{engine_args.data_parallel_degree}_cfg{engine_config.parallel_config.cfg_degree}_"
        f"ulysses{engine_args.ulysses_degree}_ring{engine_args.ring_degree}_"
        f"tp{engine_args.tensor_parallel_degree}_"
        f"pp{engine_args.pipefusion_parallel_degree}_patch{engine_args.num_pipeline_patch}"
    )

    def run_generation(current_prompts, image_offset: int):
        normalized_prompts = _normalize_prompts(current_prompts)
        current_batch_size = len(normalized_prompts) or 1
        ensure_prepared(current_batch_size)
        input_config.prompt = current_prompts
        start_time = torch.cuda.Event(enable_timing=True)
        end_time = torch.cuda.Event(enable_timing=True)
        start_time.record()
        output = pipe(
            height=input_config.height,
            width=input_config.width,
            prompt=current_prompts,
            num_inference_steps=input_config.num_inference_steps,
            output_type=input_config.output_type,
            max_sequence_length=256,
            guidance_scale=0.0,
            joint_attention_kwargs=joint_attention_kwargs,
            generator=torch.Generator(device="cuda").manual_seed(input_config.seed),
        )

        end_time.record()
        torch.cuda.synchronize()
        elapsed_time = start_time.elapsed_time(end_time) * 1e-3
        peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}")

        if input_config.output_type == "pil":
            dp_group_index = get_data_parallel_rank()
            num_dp_groups = get_data_parallel_world_size()
            dp_batch_size = (input_config.batch_size + num_dp_groups - 1) // num_dp_groups
            if pipe.is_dp_last_group():
                for i, image in enumerate(output.images):
                    image_rank = dp_group_index * dp_batch_size + i + image_offset
                    image_name = (
                        f"flux_result_{parallel_info}_{image_rank}_tc_{engine_args.use_torch_compile}.png"
                    )
                    image.save(output_dir / image_name)
                    print(f"image {i} saved to {output_dir / image_name}")

        if get_world_group().rank == get_world_group().world_size - 1:
            print(
                f"epoch time: {elapsed_time:.2f} sec, parameter memory: {parameter_peak_memory/1e9:.2f} GB, memory: {peak_memory/1e9:.2f} GB"
            )

    processed_prompts = 0
    if prompt_batches:
        image_offset = 0
        for batch_prompts in prompt_batches:
            run_generation(batch_prompts, image_offset)
            image_offset += len(batch_prompts)
            processed_prompts += len(batch_prompts)
            if processed_prompts >= total_prompts:
                break
    else:
        run_generation(input_config.prompt, 0)
    get_runtime_state().destroy_distributed_env()


if __name__ == "__main__":
    main()

set -x

export PYTHONPATH=$PWD:$PYTHONPATH

# Select the model type
#export MODEL_TYPE="FLUX-teacache"
export MODEL_TYPE="FLUX-TaylorSeer"
# Configuration for different model types
# script, model_id, inference_step
declare -A MODEL_CONFIGS=(
    # ["Pixart-alpha"]="pixartalpha_example.py /cfs/dit/PixArt-XL-2-1024-MS 20"
    # ["Pixart-sigma"]="pixartsigma_example.py /cfs/dit/PixArt-Sigma-XL-2-2K-MS 20"
    # ["Sd3"]="sd3_example.py /cfs/dit/stable-diffusion-3-medium-diffusers 20"
    ["FLUX-TaylorSeer"]="./xfusers_taylorseer_flux.py /home/MODELS/FLUX.1-dev 50"
    # ["Flux"]="flux_example.py black-forest-labs/FLUX.1-dev 28"
    # ["HunyuanDiT"]="hunyuandit_example.py /cfs/dit/HunyuanDiT-v1.2-Diffusers 50"
)

if [[ -v MODEL_CONFIGS[$MODEL_TYPE] ]]; then
    IFS=' ' read -r SCRIPT MODEL_ID INFERENCE_STEP <<< "${MODEL_CONFIGS[$MODEL_TYPE]}"
    export SCRIPT MODEL_ID INFERENCE_STEP
else
    echo "Invalid MODEL_TYPE: $MODEL_TYPE"
    exit 1
fi

OUTPUT_DIR=${OUTPUT_DIR:-./results}
PROMPT=${PROMPT:-"A dog standing on the grass."}
PROMPT_FILE=${PROMPT_FILE:-}
PROMPT_FILE_BATCH_SIZE=${PROMPT_FILE_BATCH_SIZE:-}

mkdir -p "${OUTPUT_DIR}"

# task args
TASK_ARGS="--height 1024 --width 1024 --no_use_resolution_binning"

# pipefusion is not recommended for cache based methods.
# because the computation is greatly reduced,
# and the communication cost cannot be hidden asynchronously.

export CUDA_VISIBLE_DEVICES=0,1
N_GPUS=2
PARALLEL_ARGS="--pipefusion_parallel_degree 1 --ulysses_degree 1 --ring_degree ${N_GPUS}"

# CFG_ARGS="--use_cfg_parallel"

# By default, num_pipeline_patch = pipefusion_degree, and you can tune this parameter to achieve optimal performance.
#PIPEFUSION_ARGS="--num_pipeline_patch 8 "

# For high-resolution images, we use the latent output type to avoid runing the vae module. Used for measuring speed.
# OUTPUT_ARGS="--output_type latent"

# PARALLLEL_VAE="--use_parallel_vae"

# Another compile option is `--use_onediff` which will use onediff's compiler.
# COMPILE_FLAG="--use_torch_compile"


# Use this flag to quantize the T5 text encoder, which could reduce the memory usage and have no effect on the result quality.
# QUANTIZE_FLAG="--use_fp8_t5_encoder"

# export CUDA_VISIBLE_DEVICES=4,5,6,7

CMD=(
    torchrun --nproc_per_node=$N_GPUS $SCRIPT
    --model $MODEL_ID
    $PARALLEL_ARGS
    $TASK_ARGS
    $OUTPUT_ARGS
    --num_inference_steps $INFERENCE_STEP
    --warmup_steps 1
    --prompt "$PROMPT"
    --output-dir "$OUTPUT_DIR"
)

if [[ -n "$PROMPT_FILE" ]]; then
    CMD+=(--prompt-file "$PROMPT_FILE")
fi

if [[ -n "$PROMPT_FILE_BATCH_SIZE" ]]; then
    CMD+=(--prompt-file-batch-size "$PROMPT_FILE_BATCH_SIZE")
fi

printf "  %q" "${CMD[@]}"
echo
"${CMD[@]}"


# $CFG_ARGS \
# $PARALLLEL_VAE \
# $COMPILE_FLAG \
# $QUANTIZE_FLAG \
# $PIPEFUSION_ARGS \

i don‘t know if there's any bug in this script？ the front basically means i can only get maybe three or four pictures that are of high quality, and then the pictures generated are either blurry or the image information is incomplete.

[Shenyi-Z]

This is likely due to your modifications, since we haven't seen this issue when running multi-image inference with our scripts. Also, there doesn't seem to be a need to force the model to complete multiple images in one inference pass (under 1024*1024 resolution). Based on my experiments, batch size doesn't have a noticeable impact on the speed of generating the same number of images, latency is mainly tied to total compute.

I suggest you try inferring your COCO prompts directly with the script we provided here. It worked fine for us when we ran it on DrawBench200 and Partiprompt1632.