Update app.py

app.py

@@ -1,88 +1,107 @@
 import gradio as gr
-import torch
-from PIL import Image
 
-# Список моделей
-models = [
-    "CLIP",
+from transformers import AutoModelForConditionalGeneration, AutoTokenizer
+
+# Базовые настройки
+
+prompt = gr.inputs.Textbox(label="Промпт")
+model_list = [
     "DALL-E 2",
-    "VQGAN + CLIP",
+    "VQGAN+CLIP",
+    "BigGAN",
+    "StyleGAN2",
     "VQGAN",
-    "Dreams",
-    "VQGAN-CLIP-ADA",
-    "DALL-E Mini",
-    "VQGAN-CLIP-Ada-Mini",
-    "DALL-E Mini-Fusion",
-    "VQGAN-CLIP-Ada-Mini-Fusion",
-    "DALL-E 2 Mini",
-    "VQGAN-CLIP-Ada-Mini-2",
+    "CLIP",
+    "VQGAN+CLIP-Vanilla",
+    "VQGAN+CLIP-Cutout",
+    "VQGAN+CLIP-RandomizedCutout",
 ]
+model = gr.widgets.ToggleButtons(options=model_list, label="Модель")
+
+# Расширенные настройки
+
+negative_prompt = gr.inputs.Textbox(label="Отрицательный промпт")
+sampling_method = gr.inputs.RadioButtons(options=["greedy", "top-k", "nucleus"], label="Метод выборки")
+sampling_steps = gr.inputs.Number(min=1, max=100, label="Количество шагов выборки")
+cfg_scale = gr.inputs.Number(min=0.1, max=1.0, label="Масштаб CFG")
+seed = gr.inputs.Number(min=0, max=2**31 - 1, label="Случайное число")
+
+# Улучшение качества
+
+upscale_algorithm = gr.inputs.RadioButtons(options=["bicubic", "lanczos"], label="Алгоритм увеличения")
 
 # Функция генерации изображения
-def generate_image(prompt, negative_prompt, model):
-    if model == "CLIP":
-        from clip import clip
-        image = clip.generate(prompt, negative_prompt)
-    elif model == "DALL-E 2":
-        from d2l import d2l
-        image = d2l.generate_image(prompt, negative_prompt)
-    elif model == "VQGAN + CLIP":
-        from vqgan_clip import vqgan_clip
-        image = vqgan_clip.generate(prompt, negative_prompt)
+
+def generate_image(prompt, model, negative_prompt, sampling_method, sampling_steps, cfg_scale, seed):
+    if model == "DALL-E 2":
+        model = AutoModelForConditionalGeneration.from_pretrained("google/dalle-2-1024")
+        tokenizer = AutoTokenizer.from_pretrained("google/dalle-2-1024")
+    elif model == "VQGAN+CLIP":
+        model = AutoModelForConditionalGeneration.from_pretrained("openai/vqgan-clip")
+        tokenizer = AutoTokenizer.from_pretrained("openai/vqgan-clip")
+    elif model == "BigGAN":
+        model = AutoModelForConditionalGeneration.from_pretrained("karras2022/biggan-deep-256")
+        tokenizer = AutoTokenizer.from_pretrained("karras2022/biggan-deep-256")
+    elif model == "StyleGAN2":
+        model = AutoModelForConditionalGeneration.from_pretrained("NVlabs/stylegan2-ada")
+        tokenizer = AutoTokenizer.from_pretrained("NVlabs/stylegan2-ada")
     elif model == "VQGAN":
-        from vqgan import vqgan
-        image = vqgan.generate(prompt, negative_prompt)
-    elif model == "Dreams":
-        from dreams import dreams
-        image = dreams.generate(prompt, negative_prompt)
-    elif model == "VQGAN-CLIP-ADA":
-        from vqgan_clip_ada import vqgan_clip_ada
-        image = vqgan_clip_ada.generate(prompt, negative_prompt)
-    elif model == "DALL-E Mini":
-        from dall_e_mini import dall_e_mini
-        image = dall_e_mini.generate(prompt, negative_prompt)
-    elif model == "VQGAN-CLIP-Ada-Mini":
-        from vqgan_clip_ada_mini import vqgan_clip_ada_mini
-        image = vqgan_clip_ada_mini.generate(prompt, negative_prompt)
-    elif model == "DALL-E Mini-Fusion":
-        from dall_e_mini_fusion import dall_e_mini_fusion
-        image = dall_e_mini_fusion.generate(prompt, negative_prompt)
-    elif model == "VQGAN-CLIP-Ada-Mini-Fusion":
-        from vqgan_clip_ada_mini_fusion import vqgan_clip_ada_mini_fusion
-        image = vqgan_clip_ada_mini_fusion.generate(prompt, negative_prompt)
-    elif model == "DALL-E 2 Mini":
-        from dall_e_2_mini import dall_e_2_mini
-        image = dall_e_2_mini.generate(prompt, negative_prompt)
-    elif model == "VQGAN-CLIP-Ada-Mini-2":
-        from vqgan_clip_ada_mini_2 import vqgan_clip_ada_mini_2
-        image = vqgan_clip_ada_mini_2.generate(prompt, negative_prompt)
-
-    image = Image.fromarray(image)
+        model = AutoModelForConditionalGeneration.from_pretrained("vqgan/vqgan_imagenet_f16_1024")
+        tokenizer = AutoTokenizer.from_pretrained("vqgan/vqgan_imagenet_f16_1024")
+    elif model == "CLIP":
+        model = AutoModelForConditionalGeneration.from_pretrained("openai/clip")
+        tokenizer = AutoTokenizer.from_pretrained("openai/clip")
+    elif model == "VQGAN+CLIP-Vanilla":
+        model = AutoModelForConditionalGeneration.from_pretrained("vqgan/vqgan_imagenet_f16_1024_clip_vanilla")
+        tokenizer = AutoTokenizer.from_pretrained("vqgan/vqgan_imagenet_f16_1024_clip_vanilla")
+    elif model == "VQGAN+CLIP-Cutout":
+        model = AutoModelForConditionalGeneration.from_pretrained("vqgan/vqgan_imagenet_f16_1024_clip_cutout")
+        tokenizer = AutoTokenizer.from_pretrained("vqgan/vqgan_imagenet_f16_1024_clip_cutout")
+    elif model == "VQGAN+CLIP-RandomizedCutout":
+    model = AutoModelForConditionalGeneration.from_pretrained("vqgan/vqgan_imagenet_f16_1024_clip_randomized_cutout")
+    tokenizer = AutoTokenizer.from_pretrained("vqgan/vqgan_imagenet_f16_1024_clip_randomized_cutout")
+
+    prompt = f"{prompt} {negative_prompt}"
+
+    image = model.generate(
+        text=prompt,
+        sampling_method=sampling_method,
+        sampling_steps=sampling_steps,
+        cfg_scale=cfg_scale,
+        seed=seed,
+    )
+
     return image
 
 # Функция улучшения качества изображения
-def improve_image(image):
-    from improve_image import improve_image
-    image = improve_image(image)
+
+def upscale_image(image, upscale_algorithm):
+    if upscale_algorithm == "bicubic":
+        image = cv2.resize(image, dsize=(image.shape[1] * 2, image.shape[0] * 2), interpolation=cv2.INTER_CUBIC)
+    elif upscale_algorithm == "lanczos":
+        image = cv2.resize(image, dsize=(image.shape[1] * 2, image.shape[0] * 2), interpolation=cv2.INTER_LANCZOS4)
     return image
 
-# Интерфейс gradio
-app = gr.Interface(
-    generate_image,
-    title="Gradio Image Generator",
-    description="Генератор изображений с помощью ИИ",
-    inputs=[
-        gr.Textbox(label="Текстовый запрос"),
-        gr.Textbox(label="Отрицательный текстовый запрос"),
-        gr.Dropdown(choices=models, label="Модель"),
-    ],
-    outputs=[
-        gr.Image(label="Сгенерированное изображение")
-    ],
-    buttons=[
-        gr.Button("Улучшить качество", fn=improve_image)
-    ]
-)
+# Функция отображения изображения
+
+def show_image(image):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    cv2.imshow("Image", image)
+    cv2.waitKey(0)
+
+# Основная функция
+
+def main():
+    image = generate_image(prompt.value, model.value, negative_prompt.value, sampling_method.value, sampling_steps.value, cfg_scale.value, seed.value)
+
+    # Если выбрано улучшение качества изображения
+
+    if upscale_algorithm.value != "none":
+        image = upscale_image(image, upscale_algorithm.value)
+
+    # Отображение изображения
+
+    show_image(image)
 
 if __name__ == "__main__":
-    app.launch()
+    main()