inference_batch.py

import os
import copy
import torch
from PIL import Image
from tqdm import tqdm
from models.pipe import CoTylePipeline, PiCoTylePipeline

from io import BytesIO
import requests
from models.vlm_unitok import UniTok
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
import argparse
from models.utils import set_seed, load_and_process_config, patched_from_model_config, concatenate_images_with_sref
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from models.model import StyleGenerator
import json
from models.model import Qwen2_5_VLForConditionalGeneration_Quant, Qwen2_5_VL_Quant
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
from transformers.generation.configuration_utils import GenerationConfig
_original_from_model_config = GenerationConfig.from_model_config
GenerationConfig.from_model_config = classmethod(patched_from_model_config)

def main(args):
    output_dir = args.output_path

    unitok_config = {
        'unitok_embed_dim' : 3584,
        'unitok_vocab_width' : 64,
        'unitok_vocab_size' : 1024,
        'unitok_e_temp' : 0.01,
        'unitok_num_codebooks' : 1,
        'unitok_le' : 0.0
    }
    weight_type = torch.bfloat16

    style_generator_path = os.path.join(args.model_path, 'prior')
    config = AutoConfig.from_pretrained(f"{style_generator_path}/config.json")
    style_generator = StyleGenerator._from_config(config)
    state_dict = torch.load(f"{style_generator_path}/prior.pth", map_location='cpu')
    style_generator.load_state_dict(state_dict)
    style_generator.to('cuda', dtype=weight_type)
    
    
    # loading codebook
    unitok = UniTok(unitok_config)
    unitok_state_dict = torch.load(f"{args.model_path}/codebook/model.pth", map_location='cpu')
    unitok.load_state_dict(unitok_state_dict)
    unitok.to('cuda', dtype=weight_type)

    # loading text_encoder
    if args.accelerate:
        pipeline = PiCoTylePipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16, text_encoder=None,processor=None)
    else:
        pipeline = CoTylePipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16, text_encoder=None,processor=None)

    qwen_text_visual_encoder = Qwen2_5_VLForConditionalGeneration_Quant.from_pretrained(
        os.path.join(args.model_path, 'text_encoder'),
    ).to('cuda', dtype=weight_type)

    qwen_text_visual_encoder = Qwen2_5_VL_Quant(unitok, qwen_text_visual_encoder)
    qwen_text_visual_encoder.to('cuda', dtype=weight_type)

    pipeline.text_encoder = qwen_text_visual_encoder
    processor = Qwen2VLProcessor.from_pretrained(os.path.join(args.model_path, 'processor'),
                                                                min_pixels=64 * 28 * 28,
                                                                max_pixels=256 * 28 * 28)
    pipeline.processor = processor

    if args.accelerate:
        adapter_name = pipeline.load_piflow_adapter(  # you may later call `pipe.set_adapters([adapter_name, ...])` to combine other adapters (e.g., style LoRAs)
            'Lakonik/pi-Qwen-Image',
            subfolder='gmqwen_k8_piid_4step',
            target_module_name='transformer')
        pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(  # use fixed shift=3.2
            pipeline.scheduler.config, shift=3.2, shift_terminal=None, use_dynamic_shifting=False)


    pipeline.to('cuda', dtype=torch.bfloat16)

    pipeline.set_progress_bar_config(disable=True)
    os.makedirs(output_dir, exist_ok=True)
    placeholder_image = Image.new("RGB", (392, 392), (0, 0, 0)) 
    with open(f'{args.model_path}/freq.json', 'r') as f:
        code_freq = json.load(f)
    
    prompts = []
    with open(args.prompt_file_path, 'r') as f:
        for line in f:
            prompts.append(line.strip())
    
    res_imgs = []
    seeds = []
    for style_code in tqdm(args.style_code):
        set_seed(style_code)
        style_generator_inputs = dict()
        style_generator_inputs['input_ids'] = torch.randint(low=0, high=1024, size=(1, 1)).to('cuda')
        style_generator_inputs['attention_mask'] = torch.ones(style_generator_inputs['input_ids'].shape).to('cuda')
        
        generated_ids = style_generator.generate(
            **style_generator_inputs,
            max_new_tokens=195,
            temperature=1.0,  
            top_k=200,        
            top_p=0.95,      
            do_sample=True , 
            repetition_penalty=50.0,
            code_freq=code_freq,
            code_freq_threshold=args.freq_threshold,
            k=args.freq_k,
            )

        set_seed(args.seed)

        if args.accelerate:
            sample_steps = 4
        else:
            sample_steps = 40

        res_imgs_line = []
        for prompt_i, prompt in enumerate(prompts):
            
            
            inputs = {
                "image": [placeholder_image],
                "prompt": prompt,
                "generator": torch.manual_seed(args.seed),
                "true_cfg_scale": 6.0,
                "negative_prompt": "丑陋，怪物，怪兽，畸形，变异，结构不合理，肢体不合理，人脸扭曲, 肢体错乱,突兀",
                "num_inference_steps": sample_steps,
                "guidance_scale": 1.0,
                "num_images_per_prompt": 1,
                "codebook_id": generated_ids,
            }

            with torch.inference_mode():
                output = pipeline(**inputs)
            res_imgs_line.append(output.images[0])
            os.makedirs(f"{output_dir}/batch/", exist_ok=True)
            output.images[0].save(f"{output_dir}/batch/{style_code}-{prompt_i}.png")
        res_imgs.append(copy.deepcopy(res_imgs_line))
        seeds.append(style_code)

    concatenate_images_with_sref(res_imgs, seeds).save(f"{output_dir}/batch.png")
    print(f"The results are saved to {output_dir}/batch.png")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple example of a training script.")

    parser.add_argument(
        "--style_code",
        type=int,
        nargs='+',  
        default=[1234567, 5201314,13415926, 886, 20010627, 996007, 2333],
        )
    parser.add_argument(
        "--seed",
        type=int,
        default=1024,
    )

    parser.add_argument(
        "--model_path",
        type=str,
        default='./pretrained_models',
    )


    parser.add_argument(
        "--output_path",
        type=str,
        default="outputs",
    )


    parser.add_argument(
        "--prompt_file_path",
        type=str,
        default="./test_prompts.txt",
    )
    parser.add_argument(
        "--freq_threshold",
        type=int,
        default=90000,
    )
    parser.add_argument(
        "--freq_k",
        type=float,
        default=0.0001,
    )

    parser.add_argument(
        "--accelerate",
        action='store_true'
    )

    args = parser.parse_args()
    main(args)