import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import VLChatProcessor
from Upsample import RealESRGAN

def load_vl_model():
    model_path = "deepseek-ai/Janus-Pro-7B"
    config = AutoConfig.from_pretrained(model_path)
    language_config = config.language_config
    language_config._attn_implementation = 'eager'
    vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                                 language_config=language_config,
                                                 trust_remote_code=True)
    if torch.cuda.is_available():
        vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
    else:
        vl_gpt = vl_gpt.to(torch.float16)

    vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
    tokenizer = vl_chat_processor.tokenizer
    cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'

    return vl_gpt, vl_chat_processor, tokenizer, cuda_device

def load_sr_model():
    sr_model = RealESRGAN(torch.device('cuda' if torch.cuda.is_available() else 'cpu'), scale=2)
    sr_model.load_weights(f'weights/RealESRGAN_x2.pth', download=False)
    return sr_model