extract_llava_text_encoder.py

# SPDX-License-Identifier: Apache-2.0
"""
Extract LlamaModel text encoder from LLaVA weights for GameCraft/HunyuanVideo.

LLaVA stores the text encoder as `language_model.model.XXX`, but FastVideo/HunyuanVideo
expects just `XXX` (e.g., `layers.0.self_attn.q_proj.weight` instead of 
`language_model.model.layers.0.self_attn.q_proj.weight`).

This script extracts just the LlamaModel weights and renames them.

Usage:
    python scripts/checkpoint_conversion/extract_llava_text_encoder.py \
        --input Hunyuan-GameCraft-1.0/weights/stdmodels/llava-llama-3-8b-v1_1-transformers \
        --output official_weights/hunyuan-gamecraft-diffusers/text_encoder
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from collections import OrderedDict

import torch
from safetensors import safe_open
from safetensors.torch import save_file


def extract_text_encoder(
    input_dir: Path,
    output_dir: Path,
) -> dict:
    """Extract LlamaModel from LLaVA weights."""
    
    # Find all safetensor shards
    shard_files = sorted(input_dir.glob("model-*.safetensors"))
    if not shard_files:
        raise FileNotFoundError(f"No safetensor files found in {input_dir}")
    
    print(f"Found {len(shard_files)} weight shards")
    
    # Collect all weights, filtering for language_model.model prefix
    extracted_weights: OrderedDict[str, torch.Tensor] = OrderedDict()
    
    for shard_file in shard_files:
        print(f"Processing {shard_file.name}...")
        with safe_open(str(shard_file), framework="pt") as f:
            for key in f.keys():
                # Extract only language_model.model weights
                if key.startswith("language_model.model."):
                    # Strip the prefix
                    new_key = key.replace("language_model.model.", "")
                    extracted_weights[new_key] = f.get_tensor(key)
    
    print(f"Extracted {len(extracted_weights)} weights")
    
    # Save to output directory
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Save as single file (it should be <10GB which is fine for safetensors)
    output_file = output_dir / "model.safetensors"
    print(f"Saving to {output_file}...")
    save_file(extracted_weights, str(output_file))
    
    # Create config.json for LlamaModel
    config = {
        "_name_or_path": "llava-llama-3-8b-v1_1-text_encoder",
        "architectures": ["LlamaModel"],
        "attention_bias": False,
        "attention_dropout": 0.0,
        "bos_token_id": 128000,
        "eos_token_id": 128001,
        "head_dim": 128,
        "hidden_act": "silu",
        "hidden_size": 4096,
        "initializer_range": 0.02,
        "intermediate_size": 14336,
        "max_position_embeddings": 8192,
        "mlp_bias": False,
        "model_type": "llama",
        "num_attention_heads": 32,
        "num_hidden_layers": 32,
        "num_key_value_heads": 8,
        "pretraining_tp": 1,
        "rms_norm_eps": 1e-05,
        "rope_scaling": None,
        "rope_theta": 500000.0,
        "tie_word_embeddings": False,
        "torch_dtype": "float16",
        "transformers_version": "4.48.0",
        "use_cache": True,
        "vocab_size": 128320
    }
    
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)
    print(f"Saved config.json")
    
    # Copy tokenizer files if they exist
    tokenizer_files = [
        "tokenizer.json",
        "tokenizer_config.json", 
        "special_tokens_map.json",
    ]
    for fname in tokenizer_files:
        src = input_dir / fname
        if src.exists():
            import shutil
            shutil.copy2(src, output_dir / fname)
            print(f"Copied {fname}")
    
    return {"total": len(extracted_weights)}


def extract_clip_text_encoder(
    input_dir: Path,
    output_dir: Path,
) -> dict:
    """Extract CLIPTextModel from full CLIP model weights."""
    
    # Find weight file
    weight_files = list(input_dir.glob("model.safetensors")) + \
                   list(input_dir.glob("pytorch_model.bin"))
    
    if not weight_files:
        raise FileNotFoundError(f"No weight files found in {input_dir}")
    
    weight_file = weight_files[0]
    print(f"Loading {weight_file}...")
    
    if weight_file.suffix == ".safetensors":
        from safetensors.torch import load_file
        state_dict = load_file(str(weight_file))
    else:
        state_dict = torch.load(weight_file, map_location="cpu")
    
    # Extract only text_model weights
    extracted_weights: OrderedDict[str, torch.Tensor] = OrderedDict()
    
    for key, value in state_dict.items():
        # CLIP full model has text_model.XXX, we want just XXX
        if key.startswith("text_model."):
            new_key = key.replace("text_model.", "")
            extracted_weights[new_key] = value
    
    if not extracted_weights:
        # Maybe it's already just the text model
        extracted_weights = OrderedDict(state_dict)
    
    print(f"Extracted {len(extracted_weights)} weights")
    
    # Save to output directory  
    output_dir.mkdir(parents=True, exist_ok=True)
    
    output_file = output_dir / "model.safetensors"
    print(f"Saving to {output_file}...")
    save_file(extracted_weights, str(output_file))
    
    # Create config.json for CLIPTextModel
    config = {
        "_name_or_path": "openai/clip-vit-large-patch14",
        "architectures": ["CLIPTextModel"],
        "attention_dropout": 0.0,
        "bos_token_id": 0,
        "dropout": 0.0,
        "eos_token_id": 2,
        "hidden_act": "quick_gelu",
        "hidden_size": 768,
        "initializer_factor": 1.0,
        "initializer_range": 0.02,
        "intermediate_size": 3072,
        "layer_norm_eps": 1e-05,
        "max_position_embeddings": 77,
        "model_type": "clip_text_model",
        "num_attention_heads": 12,
        "num_hidden_layers": 12,
        "pad_token_id": 1,
        "projection_dim": 768,
        "torch_dtype": "float16",
        "transformers_version": "4.48.0",
        "vocab_size": 49408
    }
    
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)
    print(f"Saved config.json")
    
    # Copy tokenizer files if they exist
    tokenizer_files = [
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
        "vocab.json",
        "merges.txt",
    ]
    for fname in tokenizer_files:
        src = input_dir / fname
        if src.exists():
            import shutil
            shutil.copy2(src, output_dir / fname)
            print(f"Copied {fname}")
    
    return {"total": len(extracted_weights)}


def main():
    parser = argparse.ArgumentParser(
        description="Extract text encoder from LLaVA or CLIP model weights."
    )
    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="Input directory containing model weights"
    )
    parser.add_argument(
        "--output",
        type=str,
        required=True,
        help="Output directory for extracted weights"
    )
    parser.add_argument(
        "--type",
        type=str,
        choices=["llama", "clip"],
        default="llama",
        help="Type of text encoder to extract"
    )
    
    args = parser.parse_args()
    
    input_dir = Path(args.input)
    output_dir = Path(args.output)
    
    if not input_dir.exists():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")
    
    if args.type == "llama":
        stats = extract_text_encoder(input_dir, output_dir)
    else:
        stats = extract_clip_text_encoder(input_dir, output_dir)
    
    print(f"\nExtraction complete! Total weights: {stats['total']}")


if __name__ == "__main__":
    main()