Spaces:
Running on Zero
Running on Zero
File size: 1,338 Bytes
3bd5145 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM
class Qwen3TextEncoder(nn.Module):
"""Wraps Qwen3 (decoder-only) as a text encoder for Wan audio pipeline.
Loads the full Qwen3 model and extracts last-layer hidden states
as text embeddings. Interface matches WanTextEncoder.forward(ids, mask).
"""
def __init__(self, model_path, torch_dtype=torch.bfloat16):
super().__init__()
self.model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch_dtype,
output_hidden_states=True,
)
self.model.eval()
for param in self.model.parameters():
param.requires_grad = False
self.dim = self.model.config.hidden_size # 2048 for Qwen3-1.7B
def forward(self, ids, mask=None):
"""
Args:
ids: [batch, seq_len] token ids
mask: [batch, seq_len] attention mask (1=valid, 0=pad)
Returns:
hidden_states: [batch, seq_len, dim] last-layer hidden states
"""
with torch.no_grad():
outputs = self.model(
input_ids=ids,
attention_mask=mask,
output_hidden_states=True,
use_cache=False,
)
return outputs.hidden_states[-1]
|