init

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 ivanovot
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

QuckDrawGAN/__init__.py

@@ -0,0 +1,2 @@
+from .train import train, discriminator_fine_tune
+from .model import Model

QuckDrawGAN/model.py

@@ -0,0 +1,89 @@
+import argparse
+import torch
+import numpy as np
+import torchvision.utils as vutils
+from .utils.models import Generator, Discriminator, latent_dim
+import hashlib
+from PIL import Image
+import warnings
+warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+
+class Model:
+    def __init__(self, generator_path, discriminator_path=None):
+        # Определяем устройство для выполнения (GPU или CPU)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Инициализация и загрузка генератора
+        self.generator = Generator(latent_dim).to(self.device)
+        self.generator.load_state_dict(torch.load(generator_path, map_location=self.device, weights_only=True))  # Загружаем веса генератора
+        self.generator.eval()  # Переводим генератор в режим оценки
+
+        # Инициализация дискриминатора, если задан путь к его весам
+        if discriminator_path:
+            self.discriminator = Discriminator().to(self.device)
+            self.discriminator.load_state_dict(torch.load(discriminator_path, map_location=self.device, weights_only=True))  # Загружаем веса дискриминатора
+            self.discriminator.eval()  # Переводим дискриминатор в режим оценки
+        else:
+            self.discriminator = None  # Если дискриминатор не используется
+
+    def generate(self, n=1, seed=None):
+        """Генерирует n изображений. Если дискриминатор загружен, возвращает изображение с наибольшей оценкой дискриминатора."""
+        with torch.no_grad():  # Отключаем градиенты для режима оценки
+            # Установка сида для воспроизводимости, если задан
+            if seed is not None:
+                seed_number = int(hashlib.md5(seed.encode()).hexdigest(), 16) % (2**32)  # Преобразуем текстовый сид в число
+                torch.manual_seed(seed_number)  # Устанавливаем сид для генерации
+
+            # Генерация случайного латентного вектора
+            z = torch.randn(n, latent_dim).to(self.device)
+
+            # Генерация изображений
+            gen_imgs = self.generator(z)
+
+            # Если дискриминатор загружен, выбираем изображение с наилучшей оценкой
+            if self.discriminator:
+                predictions = self.discriminator(gen_imgs).cpu().numpy().flatten()  # Получаем оценки дискриминатора
+                max_pred_idx = predictions.argmax()  # Находим индекс изображения с максимальной оценкой
+                best_img = gen_imgs[max_pred_idx].cpu().squeeze().numpy()  # Преобразуем изображение в формат (H, W)
+                return best_img  # Возвращаем лучшее изображение
+            else:
+                # Если дискриминатор не загружен, возвращаем первое сгенерированное изображение
+                return gen_imgs[0].cpu().squeeze().numpy()  # Преобразуем изображение в формат (H, W)
+
+if __name__ == "__main__":
+    # Определение аргументов командной строки
+    parser = argparse.ArgumentParser(description="Generate image using pretrained GAN model")
+    parser.add_argument('--generator_path', type=str, required=True, help='Path to generator model weights')
+    parser.add_argument('--discriminator_path', type=str, help='Path to discriminator model weights (optional)')
+    parser.add_argument('--output_path', type=str, default='result.png', help='Path to save the generated image')
+    parser.add_argument('--n', type=int, default=1, help='Number of images to generate')
+    parser.add_argument('--seed', type=str, help='Seed for random generation (optional)')
+
+    args = parser.parse_args()
+
+    # Инициализация модели и генерация изображения
+    model = Model(args.generator_path, args.discriminator_path)
+    generated_image = model.generate(n=args.n, seed=args.seed)
+
+    # Нормализация изображения
+    min_val = np.min(generated_image)
+    max_val = np.max(generated_image)
+
+    # Применяем нормализацию
+    normalized_image = (generated_image - min_val) / (max_val - min_val) * 255
+
+    # Приводим к 8-битному формату
+    normalized_image = normalized_image.astype(np.uint8)
+
+    # Проверяем размерность и преобразуем в RGB, если это необходимо
+    if normalized_image.ndim == 2:  # Если изображение в градациях серого
+        # Преобразуем в RGB (64, 64) -> (64, 64, 3)
+        normalized_image = np.stack([normalized_image] * 3, axis=-1)
+    elif normalized_image.shape[2] == 1:  # Если изображение с одним каналом
+        # Удаляем канал и преобразуем в RGB
+        normalized_image = np.squeeze(normalized_image, axis=2)
+
+    # Создаем изображение в формате RGB
+    img = Image.fromarray(normalized_image, mode='RGB')
+    img.save(args.output_path)

QuckDrawGAN/train.py

@@ -0,0 +1,198 @@
+import argparse
+import os
+import logging
+import random
+import torch
+import torch.optim as optim
+import torch.nn as nn
+import torchvision.utils as vutils
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from .utils.models import Generator, Discriminator, latent_dim
+from .utils.data import DrawDataset
+
+
+def train(epochs, batch_size, data_path, output_path='output', lr_g=0.001, lr_d=0.002, data_max_size=None):
+    # Создание директорий для сохранения изображений и моделей
+    os.makedirs(os.path.join(output_path, 'images'), exist_ok=True)
+    os.makedirs(os.path.join(output_path, 'models'), exist_ok=True)
+    
+    # Инициализация логирования
+    log_file = os.path.join(output_path, 'training_logs.log')
+    with open(log_file, 'w'):
+        pass  # Очищаем файл логов
+    logging.basicConfig(filename=log_file, level=logging.INFO, format='%(asctime)s - %(message)s')
+
+    # Определение устройства для обучения
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logging.info(f"Using device: {device}")
+
+    logging.info("Loading dataset")
+    dataset = DrawDataset(data_path, data_max_size)
+
+    # Инициализация генератора и дискриминатора
+    generator = Generator(latent_dim).to(device)
+    discriminator = Discriminator().to(device)
+
+    # Оптимизаторы для генератора и дискриминатора
+    optimizer_G = optim.Adam(generator.parameters(), lr=lr_g, betas=(0.9, 0.999))
+    optimizer_D = optim.Adam(discriminator.parameters(), lr=lr_d, betas=(0.9, 0.999))
+
+    # Функция потерь
+    adversarial_loss = nn.L1Loss()
+
+    # Фиксированные векторы шума для генерации изображений в каждом эпохе
+    fix_z = torch.randn(64, latent_dim).to(device)
+
+    # Загрузка данных
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
+
+    logging.info("Training started")
+
+    # Основной цикл обучения
+    for epoch in range(epochs):
+        progress_bar = tqdm(dataloader, desc=f"Epoch [{epoch+1}/{epochs}]", bar_format='{l_bar}{bar:12}{r_bar}')
+        generator.train()
+
+        for i, real_imgs in enumerate(progress_bar):
+            real_imgs = real_imgs.to(device)
+            batch_size = real_imgs.size(0)
+
+            # Создание меток для реальных и поддельных изображений
+            valid_labels = torch.full((batch_size, 1), random.uniform(0.7, 1)).to(device)
+            fake_labels = torch.full((batch_size, 1), random.uniform(0, 0.3)).to(device)
+
+            # Обновление дискриминатора
+            optimizer_D.zero_grad()
+
+            # Генерация поддельных изображений
+            z = torch.randn(batch_size, latent_dim).to(device)
+            gen_imgs = generator(z)
+
+            # Вычисление потерь для реальных и поддельных изображений
+            real_preds = discriminator(real_imgs)
+            fake_preds = discriminator(gen_imgs.detach())
+
+            loss_real = adversarial_loss(real_preds, valid_labels)
+            loss_fake = adversarial_loss(fake_preds, fake_labels)
+
+            loss_D = loss_real + loss_fake
+            loss_D.backward()
+            optimizer_D.step()
+
+            # Обновление генератора
+            optimizer_G.zero_grad()
+
+            # Генерация новых поддельных изображений
+            gen_imgs = generator(z)
+
+            # Потери генератора на основе предсказаний дискриминатора
+            fake_preds_for_gen = discriminator(gen_imgs)
+            loss_G = adversarial_loss(fake_preds_for_gen, valid_labels)
+
+            loss_G.backward()
+            optimizer_G.step()
+
+            # Обновление информации в прогресс-баре
+            progress_bar.set_postfix(Loss_D=loss_D.item(), Loss_G=loss_G.item())
+
+        # Логирование итогов эпохи
+        logging.info(f"Epoch [{epoch+1}/{epochs}], Loss_D: {loss_D.item():.4f}, Loss_G: {loss_G.item():.4f}")
+
+        # Сохранение изображений и модели
+        with torch.no_grad():
+            generator.eval()
+            gen_imgs = generator(fix_z)
+            vutils.save_image(gen_imgs.data, os.path.join(output_path, 'images', f'{epoch+1}.png'), nrow=8, normalize=True)
+            torch.save(generator.state_dict(), os.path.join(output_path, 'models', 'generator.pt'))
+            torch.save(discriminator.state_dict(), os.path.join(output_path, 'models', 'discriminator.pt'))
+
+
+def discriminator_fine_tune(generator_file, discriminator_file, data_path, batch_size=64, fine_tune_epochs=10, lr_d=0.002, data_max_size=None):
+    # Определение устройства
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logging.info(f"Using device: {device}")
+
+    # Загрузка сохранённых моделей
+    generator = Generator(latent_dim).to(device)
+    discriminator = Discriminator().to(device)
+    generator.load_state_dict(torch.load(generator_file, map_location=device, weights_only=True))
+    discriminator.load_state_dict(torch.load(discriminator_file, map_location=device, weights_only=True))
+
+    # Оптимизатор для дискриминатора
+    optimizer_D = optim.Adam(discriminator.parameters(), lr=lr_d, betas=(0.9, 0.999))
+    adversarial_loss = nn.L1Loss()
+
+    # Загрузка данных
+    dataset = DrawDataset(data_path, data_max_size)
+    fine_tune_dataloader = DataLoader(dataset, batch_size=batch_size//2, shuffle=True, pin_memory=True)
+
+    logging.info(f"Fine-tuning discriminator for {fine_tune_epochs} epochs")
+
+    for epoch in range(fine_tune_epochs):
+        progress_bar = tqdm(fine_tune_dataloader, desc=f"Fine-tuning Discriminator [{epoch+1}/{fine_tune_epochs}]", bar_format='{l_bar}{bar:12}{r_bar}')
+        discriminator.train()
+
+        for i, real_imgs in enumerate(progress_bar):
+            real_imgs = real_imgs.to(device)
+            batch_size = real_imgs.size(0)
+
+            # Создание меток для реальных и поддельных изображений
+            valid_labels = torch.full((batch_size, 1), random.uniform(0.7, 1)).to(device)
+            fake_labels = torch.full((batch_size // 2, 1), random.uniform(0, 0.3)).to(device)
+
+            # Обновление дискриминатора
+            optimizer_D.zero_grad()
+
+            # Генерация поддельных изображений
+            z = torch.randn(batch_size // 2, latent_dim).to(device)
+            gen_imgs = generator(z)
+
+            # Потери для реальных и поддельных изображений
+            real_preds = discriminator(real_imgs)
+            fake_preds = discriminator(gen_imgs.detach())
+
+            loss_real = adversarial_loss(real_preds, valid_labels)
+            loss_fake = adversarial_loss(fake_preds, fake_labels)
+
+            loss_D = loss_real + loss_fake
+            loss_D.backward()
+            optimizer_D.step()
+
+            # Обновление информации в прогресс-баре
+            progress_bar.set_postfix(Loss_D=loss_D.item())
+
+        # Логирование результатов дообучения
+        logging.info(f"Fine-tune Epoch [{epoch+1}/{fine_tune_epochs}], Loss_D: {loss_D.item():.4f}")
+
+    # Сохранение обновлённого дискриминатора
+    torch.save(discriminator.state_dict(), os.path.join(os.path.dirname(discriminator_file), 'discriminator_fine_tuned.pt'))
+
+
+# Определение аргументов командной строки
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train GAN model with specified parameters.")
+    
+    # Основные аргументы для функции train
+    parser.add_argument('--epochs', type=int, default=10, help='Number of epochs for training')
+    parser.add_argument('--batch_size', type=int, default=64, help='Batch size for training')
+    parser.add_argument('--data_path', type=str, required=True, help='Path to training data')
+    parser.add_argument('--output_path', type=str, default='output', help='Directory to save outputs')
+    parser.add_argument('--lr_g', type=float, default=0.001, help='Learning rate for generator')
+    parser.add_argument('--lr_d', type=float, default=0.002, help='Learning rate for discriminator')
+    parser.add_argument('--data_max_size', type=int, default=None, help='Maximum size of data to use')
+
+    # Аргументы для дообучения дискриминатора
+    parser.add_argument('--fine_tune', action='store_true', help='Fine-tune discriminator')
+    parser.add_argument('--generator_file', type=str, help='Path to generator weights for fine-tuning')
+    parser.add_argument('--discriminator_file', type=str, help='Path to discriminator weights for fine-tuning')
+    parser.add_argument('--fine_tune_epochs', type=int, default=10, help='Number of epochs for fine-tuning discriminator')
+    
+    args = parser.parse_args()
+    
+    if args.fine_tune:
+        # Запуск функции дообучения дискриминатора
+        discriminator_fine_tune(args.generator_file, args.discriminator_file, args.data_path, args.batch_size, args.fine_tune_epochs, args.lr_d, args.data_max_size)
+    else:
+        # Запуск основной функции тренировки
+        train(args.epochs, args.batch_size, args.data_path, args.output_path, args.lr_g, args.lr_d, args.data_max_size)

QuckDrawGAN/utils/data.py

@@ -0,0 +1,63 @@
+import torch
+from torch.utils.data import Dataset
+import pandas as pd
+from PIL import Image, ImageDraw
+import numpy as np
+
+class DrawDataset(Dataset):
+    def __init__(self, file_path, data_max_size=None):
+        # Загрузка данных из файла формата JSON
+        self.data = pd.read_json(file_path, lines=True)
+        # Отбор только распознанных рисунков
+        self.data = self.data[self.data['recognized'] == True]
+        # Ограничение размера набора данных, если указано
+        if data_max_size and len(self.data) > data_max_size:
+            self.data = self.data[:data_max_size]
+        
+        self.images = self.data['drawing'].values
+        self.processed_images = []
+
+        # Преобразование набора в изображения и нормализация
+        for raw_drawing in self.images:
+            img = self.stroke_to_image(raw_drawing)
+            img = np.array(img).astype(np.float32) / 255.0  # Нормализация изображения в диапазон [0, 1]
+            img = torch.from_numpy(img)  # Преобразование в тензор PyTorch
+            self.processed_images.append(img.unsqueeze(0))  # Добавление оси канала (1, 64, 64)
+
+    def stroke_to_image(self, raw_drawing):
+        # Коэффициенты для изменения размера изображения и его улучшения
+        scale_factor = 0.22  # Масштаб для уменьшения координат рисунков
+        upscale_factor = 8  # Коэффициент увеличения для получения плавных линий
+        original_size = 64  # Окончательный размер изображения
+        large_size = original_size * upscale_factor  # Увеличенный размер для рисования линий
+
+        # Преобразование координат линий с масштабированием и смещением
+        polylines = (
+            zip([(x + 25) * scale_factor * upscale_factor for x in polyline[0]],
+                [(y + 25) * scale_factor * upscale_factor for y in polyline[1]])
+            for polyline in raw_drawing if len(polyline) == 2
+        )
+
+        # Преобразуем набор линий в список для последующего рисования
+        polylines_list = [list(polyline) for polyline in polylines]
+        
+        # Создание пустого увеличенного изображения
+        pil_img = Image.new("L", (large_size, large_size), 255)  # Черно-белое изображение, белый фон
+        d = ImageDraw.Draw(pil_img)
+        
+        # Рисование линий с учетом масштабирования и увеличенной толщины
+        for polyline in polylines_list:
+            d.line(polyline, fill=0, width=int(1.5 * upscale_factor))  # Линии черного цвета
+        
+        # Масштабирование изображения обратно до 64x64 с использованием LANCZOS для сглаживания
+        pil_img = pil_img.resize((original_size, original_size), Image.Resampling.LANCZOS)
+        
+        return pil_img
+
+    def __len__(self):
+        # Возвращает количество изображений в наборе данных
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        # Возвращает обработанное изображение по индексу
+        return self.processed_images[idx]

QuckDrawGAN/utils/models.py

@@ -0,0 +1,107 @@
+import torch.nn as nn
+
+latent_dim = 100  # Размерность латентного пространства (входного шума для генератора)
+
+# Класс генератора для генерации изображений из латентного пространства
+class Generator(nn.Module):
+    def __init__(self, latent_dim):
+        super(Generator, self).__init__()
+
+        # Размеры начального изображения и количество каналов для начала транспонированных операций
+        self.init_size = 4  # Размер изображения после первой линейной трансформации
+        self.start_channels = 512  # Количество каналов на первом этапе
+
+        # Последовательная модель генератора
+        self.model = nn.Sequential(
+            # Линейное преобразование латентного вектора в развернутую форму для дальнейшего увеличения
+            nn.Linear(latent_dim, self.start_channels * self.init_size ** 2),
+            nn.BatchNorm1d(self.start_channels * self.init_size ** 2, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            # Преобразование в 4D тензор для начала операций с изображениями
+            nn.Unflatten(1, (self.start_channels, self.init_size, self.init_size)),
+
+            # Начало операций с изображением (увеличение размера)
+            nn.Upsample(scale_factor=2),  # Увеличение размера изображения в 2 раза
+
+            # Сверточные слои с уменьшением количества каналов и последующими нелинейностями
+            nn.Conv2d(self.start_channels, self.start_channels // 3, 3, stride=1, padding=1),
+            nn.BatchNorm2d(self.start_channels // 3, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Upsample(scale_factor=2),  # Еще одно увеличение
+
+            nn.Conv2d(self.start_channels // 3, self.start_channels // 4, 3, stride=1, padding=1),
+            nn.BatchNorm2d(self.start_channels // 4, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Upsample(scale_factor=2),  # Третье увеличение
+
+            nn.Conv2d(self.start_channels // 4, self.start_channels // 6, 3, stride=1, padding=1),
+            nn.BatchNorm2d(self.start_channels // 6, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Upsample(scale_factor=2),  # Четвертое увеличение
+
+            nn.Conv2d(self.start_channels // 6, self.start_channels // 8, 3, stride=1, padding=1),
+            nn.BatchNorm2d(self.start_channels // 8, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            # Последний сверточный слой для вывода изображения размером 1xWxH с функцией активации Тангенс
+            nn.Conv2d(self.start_channels // 8, 1, 3, stride=1, padding=1),
+            nn.Tanh()  # Приведение значений пикселей в диапазон [-1, 1]
+        )
+
+    def forward(self, z):
+        # Прямое распространение через сеть генератора
+        out = self.model(z)
+        return out  # Возвращаем сгенерированное изображение
+
+
+# Класс дискриминатора для различения реальных и сгенерированных изображений
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+
+        # Последовательная модель дискриминатора
+        self.model = nn.Sequential(
+            # Первый сверточный блок
+            nn.Conv2d(1, 64, 3, stride=2, padding=1),  # Уменьшает размер изображения до (64, 32, 32)
+            nn.BatchNorm2d(64, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout(0.3),  # Вероятность выключения нейронов для регуляризации
+
+            # Второй сверточный блок
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),  # Уменьшает размер до (128, 16, 16)
+            nn.BatchNorm2d(128, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout(0.3),
+
+            # Третий сверточный блок
+            nn.Conv2d(128, 256, 3, stride=2, padding=1),  # Уменьшает размер до (256, 8, 8)
+            nn.BatchNorm2d(256, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout(0.3),
+
+            # Четвертый сверточный блок
+            nn.Conv2d(256, 256, 3, stride=1, padding=1),  # Поддерживает размер (256, 8, 8)
+            nn.BatchNorm2d(256, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.MaxPool2d(2),  # Уменьшает размер до (256, 4, 4)
+            nn.Dropout(0.3),
+
+            # Пятый сверточный блок
+            nn.Conv2d(256, 512, 3, stride=1, padding=1),  # Уменьшает размер до (512, 2, 2)
+            nn.BatchNorm2d(512, 0.8),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.MaxPool2d(2),  # Размер до (512, 1, 1)
+
+            # Преобразование в плоский вектор для классификации
+            nn.Flatten(),  # Преобразует изображение в вектор (512 * 2 * 2 = 2048)
+            nn.Linear(512 * 2 * 2, 1)  # Полносвязный слой для получения одного скалярного выхода
+        )
+
+    def forward(self, img):
+        # Прямое распространение через сеть дискриминатора
+        out = self.model(img)
+        return out  # Возвращаем вероятность принадлежности к классу "реальное" или "сгенерированное"

app.py

@@ -0,0 +1,71 @@
+import gradio as gr
+import QuckDrawGAN as qd
+import numpy as np
+from PIL import Image
+
+# Загрузим модель
+generator_file = r'pretrained_output/models/generator.pt'
+discriminator_file = r'pretrained_output/models/discriminator_fine_tuned.pt'
+
+# Создаем объект модели
+model = qd.Model(generator_file, discriminator_file)
+
+# Функция для генерации изображения с учетом сида, нормализацией и изменением размера
+def generate_image(n_images=16, seed=""):
+    # Если сид не задан, не передаем его в модель
+    if seed == "":
+        seed = None
+    best_image = model.generate(n_images, seed)  # Генерация с учетом сида
+
+    # Нормализация: находим минимум и максимум в изображении
+    best_image_min = np.min(best_image)
+    best_image_max = np.max(best_image)
+
+    # Нормализуем изображение, чтобы значения были в диапазоне от 0 до 255
+    normalized_image = 255 * (best_image - best_image_min) / (best_image_max - best_image_min)
+
+    # Преобразуем изображение в формат, подходящий для отображения
+    pil_image = Image.fromarray(normalized_image.astype(np.uint8))  # Преобразуем в uint8 для отображения
+    pil_image = pil_image.resize((256, 256), Image.Resampling.LANCZOS)  # Ресайз изображения до 256x256
+
+    return pil_image
+
+# Создаем интерфейс Gradio с помощью Blocks (для большей гибкости)
+with gr.Blocks() as interface:
+    gr.Markdown("# Генератор изображений с использованием QuckDrawGAN")
+    gr.Markdown("Этот интерфейс позволяет генерировать изображения с помощью модели QuckDrawGAN. Настройте количество генерируемых изображений и задайте сид для повторяемости.")
+
+    with gr.Row():
+        # Блок для изображения сверху
+        generated_image = gr.Image(type="pil", label="Сгенерированное изображение", elem_id="generated_image", scale=2)  # Увеличиваем масштаб изображения
+
+    with gr.Row():
+        # Блок параметров и кнопки
+        with gr.Column():
+            seed_input = gr.Textbox(value="", label="Сид (опционально)", interactive=True)
+            num_images = gr.Slider(minimum=1, maximum=1024, value=16, label="Количество изображений для генерации", interactive=True, step=1)
+        
+        # Кнопка генерации изображения справа
+        generate_button = gr.Button("Сгенерировать")
+
+    # Логика для автогенерации при изменении параметров
+    seed_input.change(generate_image, inputs=[num_images, seed_input], outputs=generated_image)
+    num_images.change(generate_image, inputs=[num_images, seed_input], outputs=generated_image)
+
+    # Логика для кнопки генерации
+    generate_button.click(generate_image, inputs=[num_images, seed_input], outputs=generated_image)
+
+    # Автогенерация при старте
+    interface.load(generate_image, inputs=[num_images, seed_input], outputs=generated_image)
+
+    # Стилизация блока изображения (увеличение размера блока)
+    interface.css = """
+    #generated_image {
+        width: 400px;
+        height: 400px;
+        margin-top: 20px;
+    }
+    """
+
+# Запуск интерфейса
+interface.launch()