posterreward_analyser.py

#!/usr/bin/env python3
"""
PosterReward 分析模型：生成五维度文字分析
"""
import json
import os
import argparse
from swift.llm import PtEngine, InferRequest, RequestConfig


def create_analysis_prompt(prompt_text: str) -> str:
    return (
        f'Please analyze this image generated from the prompt: "{prompt_text}". '
        'Provide a detailed analysis across these five dimensions: '
        '1. Fundamental Image Integrity, '
        '2. AI Artifact and Realism Evaluation, '
        '3. Typographical Precision Analysis, '
        '4. Visual Prompt Interpretation, and '
        '5. Standalone Artistic Evaluation.'
    )


def main():
    parser = argparse.ArgumentParser(description="PosterReward 分析模型推理")
    parser.add_argument("--model", type=str, required=True)
    parser.add_argument("--prompt", type=str, required=True)
    parser.add_argument("--image_path", type=str, required=True)
    parser.add_argument("--output", type=str, default="./analysis_output.jsonl")
    parser.add_argument("--gpu", type=str, default="0")
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--max_tokens", type=int, default=2048)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    print(f"Loading analyser model: {args.model}")
    engine = PtEngine(args.model, max_batch_size=args.batch_size)

    analysis_prompt = create_analysis_prompt(args.prompt)
    messages = [{'role': 'user', 'content': f'<image>{analysis_prompt}'}]
    request = InferRequest(messages=messages, images=[args.image_path])

    request_config = RequestConfig(
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        seed=args.seed,
    )

    print("Running analysis inference...")
    resp_list = engine.infer([request], request_config=request_config)
    analysis = resp_list[0].choices[0].message.content

    os.makedirs(os.path.dirname(os.path.abspath(args.output)) or '.', exist_ok=True)

    record = {
        'image_path': args.image_path,
        'original_prompt': args.prompt,
        'analysis': analysis,
    }

    with open(args.output, 'w', encoding='utf-8') as f:
        f.write(json.dumps(record, ensure_ascii=False) + '\n')

    print(f"Analysis saved to: {args.output}")
    print(f"Analysis content:\n{analysis[:500]}...")


if __name__ == "__main__":
    main()