Create app.py

app.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+"""Stylegan-nada-ailanta.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1ysq4Y2sv7WTE0sW-n5W_HSgE28vaUDNE
+
+# Проект "CLIP-Guided Domain Adaptation of Image Generators"
+
+Данный проект представляет собой имплементацию подхода StyleGAN-NADA, предложенного в статье [StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators](https://arxiv.org/pdf/2108.00946).
+Представленный ниже функционал предназначен для визуализации реализованного проекта и включает в себя:
+- Сдвиг генератора по текстовому промпту
+- Генерация примеров
+- Генерация примеров из готовых пресетов
+- Веб-демо
+- Стилизация изображения из файла
+
+## 1. Установка
+"""
+
+# @title
+# Импорт нужных библиотек
+import os
+import sys
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import transforms
+from torchvision.utils import save_image
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Настройка устройства
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Установка библиотек
+!pip install ftfy regex tqdm
+!pip install git+https://github.com/openai/CLIP.git
+!pip install Ninja
+
+# Клонирование репозитория StyleGAN2
+!git clone https://github.com/rosinality/stylegan2-pytorch.git
+os.chdir("stylegan2-pytorch")
+
+# Скачивание весов StyleGAN
+!pip install gdown
+!gdown https://drive.google.com/uc?id=1EM87UquaoQmk17Q8d5kYIAHqu0dkYqdT
+!gdown https://drive.google.com/uc?id=1N0MZSqPRJpLfP4mFQCS14ikrVSe8vQlL
+from model import Generator
+
+# Параметры генератора
+latent_dim = 512
+f_generator = Generator(size=1024, style_dim=latent_dim, n_mlp=8).to(device)
+state_dict = torch.load('stylegan2-ffhq-config-f.pt', map_location=device)
+f_generator.load_state_dict(state_dict['g_ema'])
+f_generator.eval()
+
+g_generator = Generator(size=1024, style_dim=latent_dim, n_mlp=8).to(device)
+g_generator.load_state_dict(state_dict['g_ema'])
+
+# Загрузка модели CLIP
+import clip
+clip_model, preprocess = clip.load("ViT-B/32", device=device)
+
+latent_dim=512
+batch_size=4
+
+"""## 6. Готовые пресеты"""
+
+# @title Загрузка пресетов
+os.makedirs("/content/presets", exist_ok=True)
+
+!gdown --output /content/presets/sketch.pth https://drive.google.com/uc?id=1trcBvlz7jeBRLNeCyNVCXE4esW25GPaZ
+!gdown --output /content/presets/modigliani.pth https://drive.google.com/uc?id=1N4C-aTwxeOamZX2GeEElppsMv-ALKojL
+!gdown --output /content/presets/werewolf.pth https://drive.google.com/uc?id=1VZHEalFyEFGWIaHei98f9XPyHHvMBp6J
+
+
+# @title Генерация примеров из пресета
+# Загрузка генератора из файла
+def load_model(file_path, latent_dim=512, size=1024):
+
+    state_dicts = torch.load(file_path, map_location=device)
+
+    # Инициализация
+    trained_generator = Generator(size=size, style_dim=latent_dim, n_mlp=8).to(device)
+
+    # Загрузка весов
+    trained_generator.load_state_dict(state_dicts)
+
+    trained_generator.eval()
+
+    return trained_generator
+
+model_paths = {
+    "Photo -> Pencil Sketch": "/content/presets/sketch.pth",
+    "Photo -> Modigliani Painting": "/content/presets/modigliani.pth",
+    "Human -> Werewolf": "/content/presets/werewolf.pth"
+}
+
+
+"""## 8. Веб-демо"""
+!pip install gradio
+
+import gradio as gr
+
+# Функция обработки
+def generate(model_name):
+    model_path = model_paths[model_name]
+    g_generator = load_model(model_path)
+    images = []
+    with torch.no_grad():
+        w_optimized = f_generator.style(torch.randn(2, latent_dim).to(device))
+        w_plus = w_optimized.unsqueeze(1).repeat(1, f_generator.n_latent, 1).clone()
+
+        frozen_images = f_generator(w_plus.unsqueeze(0), input_is_latent=True)[0]
+        frozen_images = (frozen_images.clamp(-1, 1) + 1) / 2.0  # Нормализация к [0, 1]
+        frozen_images = frozen_images.permute(0, 2, 3, 1).cpu().numpy()
+        images.extend(frozen_images)
+        trained_images = g_generator(w_plus.unsqueeze(0), input_is_latent=True)[0]
+        trained_images = (trained_images.clamp(-1, 1) + 1) / 2.0  # Нормализация к [0, 1]
+        trained_images = trained_images.permute(0, 2, 3, 1).cpu().numpy()
+        images.extend(trained_images)
+    return images
+
+# Интерфейс
+iface = gr.Interface(
+    fn=generate,
+    inputs=gr.Dropdown(choices=list(model_paths.keys()), label="Выберите пресет"),
+    outputs=gr.Gallery(label="Результаты генерации", columns=2),
+    title="Выбор модели",
+    description="Выберите преобразование из списка."
+)
+
+iface.launch(debug=True)