Генератор изображений на пайтоне для телеграма

Нужна помощь срочно! Можете исправить ошибки в коде и сюда надо как-то пришить модель StyleGAN2 но можно и другую просто я не знаю как. Но впринципе можно использовать любую модель просто главное чтобы оно созавало изображения в зависимости от того что ему зададут.

Вот код:

import os
import random
import telebot
import numpy as np
from PIL import Image
import torch
from torchvision.utils import save_image
from stylegan2_pytorch import ModelLoader

# токен
TOKEN = 'YOUR_TELEGRAM_BOT_TOKEN'


# загрузка обученной модели StyleGAN2
def load_model():
# открытие файла .pkl
model_path = 'path/to/your/pretrained_model.pkl'

# загрузка
model = ModelLoader(model_path)

return model


# генерация с помощью модели
def generate_image(model):
# рандомный вектор
noise_vector = torch.randn(1, 512).cuda()

with torch.no _grad():
generated_image = model.sample(noise_vector)

return generated_image


# создание бота
bot = telebot.TeleBot(TOKEN)

# снова загрузка модели
model = load_model()


# эт старт
@bot.message_handler(commands=['start'])
def start(message):
bot.reply_to(message, "Welcome to Image Generator Bot! Type '/generate_image' to generate a new image.")


# эт обработка команды генерации
@bot.message_handler(commands=['generate_image'])
def generate(message):
# генерация
generated_image = generate_image(model)

# преобразование в PIL-изображение
generated_image_pil = Image.fromarray(np.transpose(generated_image[0].cpu().numpy(), (1, 2, 0)))

# сохранение созданной картинки
filename = 'generated_image.png'
generated_image_pil.save(filename)

# отправка
photo = open(filename, 'rb')
bot.send_photo(message.chat.id, photo)
photo.close()

# удаление изображения
os.remove(filename)

bot.polling()