gradio1.py

import os
import argparse
from pathlib import Path
import json
from typing import Optional

import torch
from PIL import Image
from transformers import AutoTokenizer

# UI
import gradio as gr
from model import MultiModalDenseTransformer
from continual_learning import UnifiedMultiModalPreprocessor

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from torchvision import transforms
image_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

class ModelInference:
    def __init__(self, checkpoint_path: str, tokenizer_name: str, config_path: Optional[str] = None, device: str = 'cuda' if torch.cuda.is_available() else 'cpu'):
        self.device = torch.device(device)
        print(f"Using device: {self.device}")
        print(f"Loading tokenizer: {tokenizer_name}...")
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True, trust_remote_code=True)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

        if config_path and Path(config_path).exists():
            with open(config_path, 'r') as f:
                self.config = json.load(f)
        else:
            self.config = {
                'model_dim': 1536,
                'vocab_size': len(self.tokenizer),
                'n_layers': 12,
                'n_heads': 12,
                'n_kv_heads': 4,
                'head_dim': None,
                'max_seq_len': 512,
                'dropout': 0.0,
                'use_moe': False,
                'use_adapter': False,
                'use_lora': False,
                'rope_scaling_type': "yarn",
                'use_multimodal_fusion': False,
                'use_contrastive': False
            }

        # init model + preprocessor
        print("Initializing model architecture...")
        self.model = MultiModalDenseTransformer(**self.config)
        self.preprocessor = UnifiedMultiModalPreprocessor(model_dim=self.config['model_dim'])

        print(f"Loading checkpoint from {checkpoint_path}...")
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        state_dict = checkpoint.get('model_state_dict', checkpoint) if isinstance(checkpoint, dict) else checkpoint

        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith('module.'):
                new_state_dict[k[7:]] = v
            else:
                new_state_dict[k] = v

        missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
        if missing:
            print(f"Warning: Missing keys: {len(missing)}")
        if unexpected:
            print(f"Warning: Unexpected keys: {len(unexpected)}")

        self.model.to(self.device)
        self.preprocessor.to(self.device)
        self.model.eval()
        print("Model loaded successfully!")
        print(f"Total parameters: {sum(p.numel() for p in self.model.parameters())/1e6:.2f}M")

    @torch.no_grad()
    def generate_text(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.7, top_k: int = 10, top_p: float = 0.9, repetition_penalty: float = 1.2, image: Optional[Image.Image] = None) -> str:
        formatted_prompt = f"Instruction: {prompt}\nResponse:"
        inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
        input_ids = inputs['input_ids'].to(self.device)

        input_data = {'segments': []}
        if image is not None:
            try:
                if image.mode != 'RGB':
                    image = image.convert('RGB')
                image_tensor = image_transform(image).unsqueeze(0).to(self.device)
                mod_segments = self.preprocessor.process_batch(image_tensor, 'image')
                for seg in mod_segments:
                    input_data['segments'].append(seg)
            except Exception as e:
                print(f"Warning: Image processing skipped due to error: {e}")

        input_data['segments'].append({
            'type': 'text',
            'data': input_ids,
            'modality_id': 0
        })

        try:
            generated_ids = self.model.generate(
                input_data,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                do_sample=True,
                eos_token_id=self.tokenizer.eos_token_id,
                pad_token_id=self.tokenizer.pad_token_id
            )

            full_output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
            # 提取 Response 后的部分并做 stop 处理
            if "Response:" in full_output:
                answer = full_output.split("Response:")[-1].strip()
            else:
                answer = full_output

            stop_words = ["Instruction", "Input", "###", "Response", "User:", "Assistant:", "\n\n"]
            for sw in stop_words:
                if sw in answer:
                    answer = answer.split(sw)[0].strip()

            # 去掉可能的 echo
            lines = answer.split('\n')
            if len(lines) > 0 and prompt.lower() in lines[0].lower():
                answer = "\n".join(lines[1:]).strip()
            return answer
        except Exception as e:
            import traceback
            traceback.print_exc()
            return f"Error: {e}"

def build_ui(model_instance):
    with gr.Blocks(title="MultiModal Dense Transformer - Gradio", css="""
        .gradio-container { max-width: 900px; margin: auto; }
    """) as demo:
        gr.Markdown("##  多模态在线推理（文本 + 图片）")
        with gr.Row():
            with gr.Column(scale=3):
                txt = gr.Textbox(label="Prompt (Instruction)", placeholder="请输入指令或问题...", lines=5)
                img = gr.Image(type="pil", label="(可选) 上传图片（支持多模态）")
                btn = gr.Button("生成 (Generate)")
            with gr.Column(scale=2):
                max_tokens = gr.Slider(label="Max New Tokens", minimum=16, maximum=1024, step=1, value=128)
                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.5, step=0.01, value=0.7)
                top_k = gr.Slider(label="Top-k", minimum=0, maximum=200, step=1, value=40)
                top_p = gr.Slider(label="Top-p", minimum=0.0, maximum=1.0, step=0.01, value=0.9)
                rep_pen = gr.Slider(label="Repetition Penalty", minimum=0.5, maximum=2.0, step=0.01, value=1.1)
                status = gr.Textbox(label="Status", value="Ready", interactive=False)
        output = gr.Textbox(label="Output", lines=12, interactive=False)

        def gr_generate(prompt, image, max_tokens_v, temp_v, topk_v, topp_v, rep_v):
            if not prompt or str(prompt).strip() == "":
                return "", "请输入 Prompt", ""
            status_msg = "Generating..."
            # call model
            out = model_instance.generate_text(prompt=prompt,
                                               max_new_tokens=int(max_tokens_v),
                                               temperature=float(temp_v),
                                               top_k=int(topk_v),
                                               top_p=float(topp_v),
                                               repetition_penalty=float(rep_v),
                                               image=image)
            return out, "Done", ""

        btn.click(fn=gr_generate, inputs=[txt, img, max_tokens, temperature, top_k, top_p, rep_pen], outputs=[output, status, gr.State()])

        demo.launch(share=True)

    return demo

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str, default="/root/multimodal/checkpoints/posttrain/final_model.pt")
    parser.add_argument("--tokenizer", type=str, default="Qwen/Qwen2.5-7B-Instruct")
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--share", type=lambda x: x.lower() in ("true","1","yes"), default=True)
    args = parser.parse_args()

    if not Path(args.checkpoint).exists():
        possible = list(Path("checkpoints/pretrain").glob("step_*.pt"))
        if possible:
            args.checkpoint = str(possible[-1])
            print(f"未找到 final_model.pt，使用最新 checkpoint: {args.checkpoint}")
        else:
            raise FileNotFoundError(f"找不到检查点: {args.checkpoint}")

    global model_instance
    model_instance = ModelInference(args.checkpoint, args.tokenizer, args.config)

    demo = build_ui(model_instance)
    demo.launch(server_port=args.port, share=args.share)

if __name__ == "__main__":
    main()