model_loader.py

from tokenxxx import *
from constants import *
from utils import *
import os
import json
import urllib.request
import urllib.parse
import torch
import hashlib
from tqdm import tqdm
from skimage import img_as_ubyte
from torch import nn
import torch.nn.functional as F
import inspect

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def filter_kwargs(cls, kwargs):
    sig = inspect.signature(cls.__init__)
    accepted = set(sig.parameters.keys()) - {"self"}
    return {k: v for k, v in kwargs.items() if k in accepted}

def sanitize_filename(name, url=None):
    for c in '<>:"/\\|?*':
        name = name.replace(c, '')
    if not name and url is not None:
        name = hashlib.md5(url.encode()).hexdigest()
    return name

def download_file(url, filepath):
    d = os.path.dirname(filepath)
    if d and not os.path.exists(d):
        os.makedirs(d, exist_ok=True)
    while not os.path.exists(filepath):
        try:
            def prog(t):
                last = [0]
                def inner(n, bs, ts):
                    if ts > 0:
                        t.total = ts
                    t.update(n * bs - last[0])
                    last[0] = n * bs
                return inner
            with tqdm(unit='B', unit_scale=True, unit_divisor=1024, desc=os.path.basename(filepath)) as t:
                urllib.request.urlretrieve(url, filepath, reporthook=prog(t))
        except Exception:
            continue

def download_files(folder, files_spec):
    if isinstance(files_spec, dict):
        for fn, url in files_spec.items():
            fn = sanitize_filename(fn, url)
            fp = os.path.join(folder, fn)
            download_file(url, fp)
    elif isinstance(files_spec, list):
        for item in files_spec:
            if isinstance(item, str):
                url = item
                parsed = urllib.parse.urlparse(url)
                fn = os.path.basename(parsed.path)
                if not fn:
                    fn = hashlib.md5(url.encode()).hexdigest()
                fn = sanitize_filename(fn, url)
            elif isinstance(item, (list, tuple)) and len(item) == 2:
                url, fn = item
                fn = sanitize_filename(fn, url)
            elif isinstance(item, dict) and "filename" in item and "url" in item:
                fn = sanitize_filename(item["filename"], item["url"])
                url = item["url"]
            else:
                raise ValueError("Invalid file specification")
            fp = os.path.join(folder, fn)
            download_file(url, fp)
    else:
        raise ValueError("files_spec must be dict or list")

def read_json(fp):
    with open(fp, 'r', encoding='utf-8') as f:
        return json.load(f)

def get_codegen_tokenizer(vocab_path, merges_path):
    with open(vocab_path, 'r', encoding='utf-8') as f:
        vocab = json.load(f)
    with open(merges_path, 'r', encoding='utf-8') as f:
        merges = f.read().splitlines()
    merge_ranks = {}
    for i, merge in enumerate(merges):
        parts = merge.strip().split()
        if len(parts) == 2:
            merge_ranks[tuple(parts)] = i
    def bpe(token):
        word = list(token)
        pairs = [(word[i], word[i+1]) for i in range(len(word)-1)]
        while True:
            candidate = None
            candidate_rank = None
            candidate_index = None
            for i, pair in enumerate(pairs):
                if pair in merge_ranks:
                    rank = merge_ranks[pair]
                    if candidate is None or rank < candidate_rank:
                        candidate = pair
                        candidate_rank = rank
                        candidate_index = i
            if candidate is None:
                break
            first, second = candidate
            new_word = []
            i = 0
            while i < len(word):
                if i < len(word) - 1 and word[i] == first and word[i+1] == second:
                    new_word.append(first + second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            word = new_word
            if len(word) == 1:
                break
            pairs = [(word[i], word[i+1]) for i in range(len(word)-1)]
        return word
    def tokenizer(text):
        tokens = []
        for token in text.split():
            bpe_tokens = bpe(token)
            for subtoken in bpe_tokens:
                tokens.append(vocab.get(subtoken, 0))
        return tokens
    return tokenizer

def simple_tokenizer(text, vocab, max_length=77):
    toks = text.split()
    ids = [vocab.get(t, 1) for t in toks]
    if len(ids) < max_length:
        ids = ids + [0] * (max_length - len(ids))
    else:
        ids = ids[:max_length]
    return torch.tensor(ids, dtype=torch.long).unsqueeze(0).to(device)

def load_state_dict_safe(model, loaded_state_dict):
    model_state = model.state_dict()
    new_state = {}
    for key, value in model_state.items():
        if key in loaded_state_dict and loaded_state_dict[key].shape == value.shape:
            new_state[key] = loaded_state_dict[key]
        else:
            new_state[key] = value
    model.load_state_dict(new_state, strict=False)

class GPT2Config:
    def __init__(self, vocab_size=50257, **kwargs):
        self.vocab_size = vocab_size
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class MBartConfig:
    def __init__(self, vocab_size=50265, **kwargs):
        self.vocab_size = vocab_size
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class CodeGenConfig:
    def __init__(self, vocab_size=50257, **kwargs):
        self.vocab_size = vocab_size
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class BartConfig:
    def __init__(self, vocab_size=50265, **kwargs):
        self.vocab_size = vocab_size
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class AutoencoderKLConfig:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class OpenLRMConfig:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class UNet2DConditionModelConfig:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class MusicGenConfig:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
    @classmethod
    def from_dict(cls, d):
        return cls(**d)

class GPT2LMHeadModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
        self.transformer = nn.TransformerEncoder(layer, num_layers=12)
        self.lm_head = nn.Linear(768, config.vocab_size)
    def forward(self, x):
        return self.lm_head(self.transformer(x))

class MBartForConditionalGeneration(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
        self.encoder = nn.TransformerEncoder(layer, num_layers=6)
        dlayer = nn.TransformerDecoderLayer(d_model=768, nhead=12)
        self.decoder = nn.TransformerDecoder(dlayer, num_layers=6)
        self.output_layer = nn.Linear(768, config.vocab_size)
    def forward(self, src, tgt):
        return self.output_layer(self.decoder(tgt, self.encoder(src)))

class CodeGenForCausalLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        d_model = getattr(config, "d_model", 1024)
        n_head = getattr(config, "n_head", 16)
        num_layers = getattr(config, "num_layers", 12)
        dlayer = nn.TransformerDecoderLayer(d_model=d_model, nhead=n_head)
        self.transformer_decoder = nn.TransformerDecoder(dlayer, num_layers=num_layers)
        self.lm_head = nn.Linear(d_model, config.vocab_size)
    def forward(self, tgt, memory=None):
        if memory is None:
            memory = torch.zeros_like(tgt)
        return self.lm_head(self.transformer_decoder(tgt, memory))

class BartForConditionalGeneration(nn.Module):
    def __init__(self, config):
        super().__init__()
        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
        self.encoder = nn.TransformerEncoder(layer, num_layers=6)
        dlayer = nn.TransformerDecoderLayer(d_model=768, nhead=12)
        self.decoder = nn.TransformerDecoder(dlayer, num_layers=6)
        self.output_layer = nn.Linear(768, config.vocab_size)
    def forward(self, src, tgt):
        return self.output_layer(self.decoder(tgt, self.encoder(src)))

class ResnetBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.norm1 = nn.GroupNorm(32, in_ch)
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.norm2 = nn.GroupNorm(32, out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.conv_shortcut = nn.Conv2d(in_ch, out_ch, 1)
    def forward(self, x):
        sc = self.conv_shortcut(x)
        h = F.silu(self.norm1(x))
        h = self.conv1(h)
        h = F.silu(self.norm2(x))
        h = self.conv2(h)
        return h + sc

class Downsample(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 3, stride=2, padding=1)
    def forward(self, x):
        return self.conv(x)

class DownBlock(nn.Module):
    def __init__(self, in_ch, out_ch, num_res):
        super().__init__()
        self.resnets = nn.ModuleList([ResnetBlock(in_ch if i == 0 else out_ch, out_ch) for i in range(num_res)])
        self.downsamplers = nn.ModuleList([Downsample(out_ch, out_ch)])
    def forward(self, x):
        for r in self.resnets:
            x = r(x)
        for ds in self.downsamplers:
            x = ds(x)
        return x

class Upsample(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.ConvTranspose2d(in_ch, out_ch, 4, stride=2, padding=1)
    def forward(self, x):
        return self.conv(x)

class UpBlock(nn.Module):
    def __init__(self, in_ch, out_ch, num_res):
        super().__init__()
        self.resnets = nn.ModuleList([ResnetBlock(in_ch if i == 0 else out_ch, out_ch) for i in range(num_res)])
        self.upsampler = Upsample(out_ch, out_ch)
    def forward(self, x):
        for r in self.resnets:
            x = r(x)
        return self.upsampler(x)

class AttentionBlock(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.norm = nn.GroupNorm(32, ch)
        self.query = nn.Conv2d(ch, ch, 1)
        self.key = nn.Conv2d(ch, ch, 1)
        self.value = nn.Conv2d(ch, ch, 1)
        self.proj_attn = nn.Conv2d(ch, ch, 1)
    def forward(self, x):
        b, c, h, w = x.shape
        xn = self.norm(x)
        q = self.query(xn).view(b, c, -1).permute(0, 2, 1)
        k = self.key(xn).view(b, c, -1)
        v = self.value(xn).view(b, c, -1).permute(0, 2, 1)
        attn = torch.softmax(torch.bmm(q, k) / (c ** 0.5), dim=-1)
        out = torch.bmm(attn, v).permute(0, 2, 1).view(b, c, h, w)
        return x + self.proj_attn(out)

class Encoder(nn.Module):
    def __init__(self, in_ch=3, base_ch=128, latent_ch=4):
        super().__init__()
        self.conv_in = nn.Conv2d(in_ch, base_ch, 3, padding=1)
        self.down_blocks = nn.ModuleList([
            DownBlock(base_ch, base_ch, 2),
            DownBlock(base_ch, base_ch * 2, 2),
            DownBlock(base_ch * 2, base_ch * 4, 2),
            DownBlock(base_ch * 4, base_ch * 4, 2)
        ])
        self.mid_block = nn.ModuleList([
            ResnetBlock(base_ch * 4, base_ch * 4),
            AttentionBlock(base_ch * 4),
            ResnetBlock(base_ch * 4, base_ch * 4)
        ])
        self.conv_norm_out = nn.GroupNorm(32, base_ch * 4)
        self.conv_out = nn.Conv2d(base_ch * 4, latent_ch * 2, 3, padding=1)
        self.quant_conv = nn.Conv2d(latent_ch * 2, latent_ch, 1)
    def forward(self, x):
        x = self.conv_in(x)
        for blk in self.down_blocks:
            x = blk(x)
        for m in self.mid_block:
            x = m(x)
        x = self.conv_norm_out(x)
        x = self.conv_out(x)
        return self.quant_conv(x)

class Decoder(nn.Module):
    def __init__(self, out_ch=3, base_ch=128, latent_ch=4):
        super().__init__()
        self.post_quant_conv = nn.Conv2d(latent_ch, latent_ch * 2, 1)
        self.conv_in = nn.Conv2d(latent_ch, base_ch * 4, 3, padding=1)
        self.mid_block = nn.ModuleList([
            ResnetBlock(base_ch * 4, base_ch * 4),
            AttentionBlock(base_ch * 4),
            ResnetBlock(base_ch * 4, base_ch * 4)
        ])
        self.up_blocks = nn.ModuleList([
            UpBlock(base_ch * 4, base_ch * 4, 3),
            UpBlock(base_ch * 4, base_ch * 2, 3),
            UpBlock(base_ch * 2, base_ch, 3),
            UpBlock(base_ch, base_ch, 3)
        ])
        self.conv_norm_out = nn.GroupNorm(32, base_ch)
        self.conv_out = nn.Conv2d(base_ch, out_ch, 3, padding=1)
    def forward(self, x):
        x = self.post_quant_conv(x)
        x = self.conv_in(x)
        for m in self.mid_block:
            x = m(x)
        for up in self.up_blocks:
            x = up(x)
        x = self.conv_norm_out(x)
        return self.conv_out(x)

class AutoencoderKL(nn.Module):
    def __init__(self, config):
        super().__init__()
        in_ch = config.get("in_channels", 3) if isinstance(config, dict) else config.__dict__.get("in_channels", 3)
        out_ch = config.get("out_channels", 3) if isinstance(config, dict) else config.__dict__.get("out_channels", 3)
        base_ch = config.get("base_channels", 128) if isinstance(config, dict) else config.__dict__.get("base_channels", 128)
        latent_ch = config.get("latent_channels", 4) if isinstance(config, dict) else config.__dict__.get("latent_channels", 4)
        self.encoder = Encoder(in_ch, base_ch, latent_ch)
        self.decoder = Decoder(out_ch, base_ch, latent_ch)
    def forward(self, x):
        return self.decoder(self.encoder(x))
    def decode(self, x):
        return self.decoder(x)

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.norm2 = nn.LayerNorm(embed_dim)
        hidden_dim = embed_dim * 4
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, embed_dim)
        )
    def forward(self, x):
        res = x
        x = self.norm1(x)
        x = x.transpose(0, 1)
        attn, _ = self.attn(x, x, x)
        x = attn.transpose(0, 1)
        x = res + x
        return x + self.mlp(self.norm2(x))

class VisionTransformer(nn.Module):
    def __init__(self, config):
        super().__init__()
        if isinstance(config, dict):
            self.img_size = config.get("img_size", 592)
            self.patch_size = config.get("patch_size", 16)
            self.embed_dim = config.get("hidden_size", 768)
            depth = config.get("depth", 12)
            num_heads = config.get("num_heads", 12)
        else:
            self.img_size = config.__dict__.get("img_size", 592)
            self.patch_size = config.__dict__.get("patch_size", 16)
            self.embed_dim = config.__dict__.get("hidden_size", 768)
            depth = config.__dict__.get("depth", 12)
            num_heads = config.__dict__.get("num_heads", 12)
        num_patches = (self.img_size // self.patch_size) ** 2
        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, self.embed_dim))
        self.patch_embed = nn.Conv2d(3, self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size)
        self.blocks = nn.ModuleList([TransformerBlock(self.embed_dim, num_heads) for _ in range(depth)])
        self.norm = nn.LayerNorm(self.embed_dim)
        self.register_tokens = nn.Parameter(torch.zeros(1, 4, self.embed_dim))
        self._init_weights()
    def _init_weights(self):
        nn.init.normal_(self.cls_token, std=0.02)
        nn.init.normal_(self.pos_embed, std=0.02)
    def forward(self, x):
        x = self.patch_embed(x)
        x = x.flatten(2).transpose(1, 2)
        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        return self.norm(x)[:, 0]

class OpenLRM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.encoder = nn.ModuleDict({"model": VisionTransformer(config)})
        hidden = config.get("hidden_size", 768) if isinstance(config, dict) else config.__dict__.get("hidden_size", 768)
        self.linear = nn.Linear(hidden, hidden)
    def forward(self, x):
        return self.linear(self.encoder["model"](x))

class VideoUNet(nn.Module):
    def __init__(self, in_ch=4, out_ch=4, features=None):
        super().__init__()
        if features is None:
            features = [64, 128, 256]
        self.encoder = nn.ModuleList()
        self.pool = nn.MaxPool3d(2, 2)
        self.decoder = nn.ModuleList()
        for f in features:
            self.encoder.append(nn.Sequential(
                nn.Conv3d(in_ch, f, 3, padding=1),
                nn.ReLU(inplace=True),
                nn.Conv3d(f, f, 3, padding=1),
                nn.ReLU(inplace=True)
            ))
            in_ch = f
        for f in reversed(features):
            self.decoder.append(nn.Sequential(
                nn.Conv3d(f * 2, f, 3, padding=1),
                nn.ReLU(inplace=True),
                nn.Conv3d(f, f, 3, padding=1),
                nn.ReLU(inplace=True)
            ))
        self.final_conv = nn.Conv3d(features[0], out_ch, 1)
    def forward(self, x, t, encoder_hidden_states):
        skips = []
        for enc in self.encoder:
            x = enc(x)
            skips.append(x)
            x = self.pool(x)
        for dec in self.decoder:
            skip = skips.pop()
            x = F.interpolate(x, scale_factor=2, mode='trilinear', align_corners=False)
            x = torch.cat([x, skip], dim=1)
            x = dec(x)
        return self.final_conv(x)

class SentimentClassifierModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.classifier = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Linear(256, 2)
        )
    def forward(self, x):
        return self.classifier(x)

class STTModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Linear(512, 768)
        )
    def forward(self, x):
        return self.net(x)

class TTSModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Linear(512, 768)
        )
    def forward(self, x):
        return self.net(x)

class MusicGenModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        layer = nn.TransformerEncoderLayer(d_model=768, nhead=12)
        self.transformer = nn.TransformerEncoder(layer, num_layers=12)
        self.linear = nn.Linear(768, 768)
    def forward(self, x):
        return self.linear(self.transformer(x))

class SimpleTextEncoder(nn.Module):
    def __init__(self, vocab_size=10000, embed_dim=768, max_length=77):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.max_length = max_length
    def forward(self, text_tokens):
        return self.embedding(text_tokens)

class DiffusionScheduler:
    def __init__(self, steps):
        self.steps = steps
        self.betas = torch.linspace(0.1, 0.001, steps=steps).to(device)
        self.alphas = 1 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)
    def step(self, noise, t, sample):
        alpha_bar = self.alpha_bars[t]
        alpha_bar_prev = self.alpha_bars[t-1] if t > 0 else torch.tensor(1.0, device=sample.device)
        x0 = (sample - torch.sqrt(1 - alpha_bar) * noise) / torch.sqrt(alpha_bar)
        new_sample = torch.sqrt(alpha_bar_prev) * x0 + torch.sqrt(1 - alpha_bar_prev) * noise
        return new_sample

class VideoOutput:
    def __init__(self, frames):
        self.frames = [img_as_ubyte(frame) for frame in frames[0]]

class VideoPipeline(nn.Module):
    def __init__(self, unet, vae, text_encoder, vocab):
        super().__init__()
        self.unet = unet
        self.vae = vae
        self.text_encoder = text_encoder
        self.vocab = vocab
    def forward(self, prompt: str, steps: int = 25, num_frames: int = 24):
        token_ids = simple_tokenizer(prompt, self.vocab)
        text_emb = self.text_encoder(token_ids)
        latent = torch.randn((1, 4, num_frames, 64, 64), device=device).half()
        sched = DiffusionScheduler(steps)
        for t in range(steps):
            noise = self.unet(latent, t, text_emb)
            latent = sched.step(noise, t, latent)
        frames = self.vae.decode(latent / 0.18215)
        frames = frames.clamp(0, 1).float().cpu().permute(0, 2, 3, 4, 1).numpy()
        return VideoOutput(frames)

def initialize_gpt2_model(folder, files):
    download_files(folder, files)
    config = GPT2Config()
    model = GPT2LMHeadModel(config).to(device)
    sd = torch.load(os.path.join(folder, sanitize_filename("gpt2-pytorch_model.bin")), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    enc = read_json(os.path.join(folder, sanitize_filename("encoder.json")))
    return model, enc

def initialize_translation_model(folder, files):
    download_files(folder, files)
    config = MBartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = MBartForConditionalGeneration(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    vp = os.path.join(folder, "vocab.json")
    if os.path.exists(vp):
        vocab = read_json(vp)
        model.tokenizer = lambda txt: [vocab.get(t, 0) for t in txt.split()]
    else:
        model.tokenizer = lambda txt: txt
    model.config.lang_code_to_id = {'en_XX': 0, 'es_XX': 1}
    return model

def initialize_codegen_model(folder, files):
    download_files(folder, files)
    config = CodeGenConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = CodeGenForCausalLM(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    tok = get_codegen_tokenizer(os.path.join(folder, "vocab.json"), os.path.join(folder, "merges.txt"))
    vocab = read_json(os.path.join(folder, "vocab.json"))
    idx2w = {v: k for k, v in vocab.items()}
    model.tokenizer = tok
    return model, tok, vocab, idx2w, vocab

def initialize_summarization_model(folder, files):
    download_files(folder, files)
    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = BartForConditionalGeneration(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    vp = os.path.join(folder, "vocab.json")
    if os.path.exists(vp):
        vocab_json = read_json(vp)
        vocab = set(vocab_json.keys())
        return model, vocab, vocab_json, {v: k for k, v in vocab_json.items()}
    return model, None, None, None

def initialize_imagegen_model(folder, files):
    download_files(folder, files)
    config = AutoencoderKLConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    vae = AutoencoderKL(config).to(device)
    sd = torch.load(os.path.join(folder, "diffusion_pytorch_model.bin"), map_location=device)
    load_state_dict_safe(vae, sd)
    vae.eval()
    return vae

def initialize_image_to_3d_model(folder, files):
    download_files(folder, files)
    config = OpenLRMConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model3d = OpenLRM(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model3d, sd)
    model3d.eval()
    return model3d

def initialize_text_to_video_model(folder, files):
    download_files(folder, files)
    unet_cfg = read_json(os.path.join(folder, "config.json"))
    unet_cfg = filter_kwargs(VideoUNet, unet_cfg)
    unet = VideoUNet(**unet_cfg).half().to(device)
    sd_unet = torch.load(os.path.join(folder, "diffusion_pytorch_model.fp16.bin"), map_location=device)
    load_state_dict_safe(unet, sd_unet)
    unet.eval()
    vae_cfg = read_json(os.path.join(folder, "config.json"))
    vae_cfg = filter_kwargs(AutoencoderKL, vae_cfg)
    vae = AutoencoderKL(vae_cfg).half().to(device)
    sd_vae = torch.load(os.path.join(folder, "diffusion_pytorch_model.bin"), map_location=device)
    load_state_dict_safe(vae, sd_vae)
    vae.eval()
    vp = os.path.join(folder, "vocab.json")
    text_vocab = read_json(vp) if os.path.exists(vp) else {}
    te_path = os.path.join(folder, "text_encoder.bin")
    if os.path.exists(te_path):
        text_encoder = SimpleTextEncoder(vocab_size=(max(text_vocab.values())+1) if text_vocab else 10000, embed_dim=768, max_length=77).to(device)
        sd_te = torch.load(te_path, map_location=device)
        load_state_dict_safe(text_encoder, sd_te)
    else:
        text_encoder = SimpleTextEncoder(vocab_size=(max(text_vocab.values())+1) if text_vocab else 10000, embed_dim=768, max_length=77).to(device)
    text_encoder.eval()
    return VideoPipeline(unet, vae, text_encoder, text_vocab)

def initialize_sentiment_model(folder, files):
    download_files(folder, files)
    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = SentimentClassifierModel(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    vp = os.path.join(folder, "vocab.json")
    if os.path.exists(vp):
        read_json(vp)
    return model

def initialize_stt_model(folder, files):
    download_files(folder, files)
    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = STTModel(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    vp = os.path.join(folder, "vocab.json")
    if os.path.exists(vp):
        read_json(vp)
    return model

def initialize_tts_model(folder, files):
    download_files(folder, files)
    config = BartConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = TTSModel(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    vp = os.path.join(folder, "vocab.json")
    if os.path.exists(vp):
        read_json(vp)
    return model

def initialize_musicgen_model(folder, files):
    download_files(folder, files)
    config = MusicGenConfig.from_dict(read_json(os.path.join(folder, "config.json")))
    model = MusicGenModel(config).to(device)
    sd = torch.load(os.path.join(folder, "pytorch_model.bin"), map_location=device)
    load_state_dict_safe(model, sd)
    model.eval()
    return model