Qwen3-0.6B-T5-xxl

from transformers import AutoTokenizer, AutoModel
import torch

# Define the model repository ID
model_id = "JusteLeo/Qwen3-0.6B-T5-xxl"

# Load the tokenizer and model
# trust_remote_code=True is required to load the custom projection head
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id, trust_remote_code=True)

# Move model to a device (e.g., GPU)
device = "cuda"
model.to(device)
model.eval()

# Create embeddings
prompts = [
    "A photorealistic portrait of a medieval knight in shiny armor.",
    "A futuristic cityscape at night, with flying cars and neon lights."
]

# Tokenize the prompts
inputs = tokenizer(prompts, padding=True, truncation=True, return_tensors="pt").to(device)

# Generate the embeddings
with torch.no_grad():
    embeddings = model(**inputs)

print("Embeddings generated successfully!")
print(f"Output shape: {embeddings.shape}")
# Expected output shape: (2, 4096)