Update app.py

app.py

@@ -1,5 +1,10 @@
 import gradio as gr
 from transformers import pipeline
+from diffusers import StableDiffusionPipeline  # Для генерации изображений
+import torch
+from PIL import Image
+import requests
+from io import BytesIO
 
 # Загружаем модели для анализа тональности, суммаризации текста, генерации подписей к изображениям, ответов на вопросы, перевода текста, определения эмоций, автодополнения кода, определения фейковых новостей, NER, классификации изображений, генерации кода и исправления кода
 sentiment_pipeline = pipeline("sentiment-analysis")
@@ -15,6 +20,30 @@ image_classification_pipeline = pipeline("image-classification", model="google/v
 code_generation_pipeline = pipeline("text-generation", model="deepseek-ai/deepseek-coder-1.3b-instruct")
 code_fix_pipeline = pipeline("text-generation", model="deepseek-ai/deepseek-coder-1.3b-instruct")
 
+# Загружаем модель для генерации изображений (Stable Diffusion)
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
+image_generation_pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1").to(device)
+
+# Загружаем модель для расширения изображений (Real-ESRGAN)
+def load_esrgan_model():
+    from realesrgan import RealESRGANer
+    model = RealESRGANer(
+        scale=4,
+        model_path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
+        model=None,
+        tile=0,
+        tile_pad=10,
+        pre_pad=0,
+        half=False,
+    )
+    return model
+
+esrgan_model = load_esrgan_model()
+
 # Функция для анализа тональности текста
 def analyze_sentiment(text):
     result = sentiment_pipeline(text)[0]
@@ -87,6 +116,19 @@ def fix_code(error, problem, solution, example):
     result = code_fix_pipeline(prompt, max_length=200, num_return_sequences=1)
     return result[0]['generated_text']
 
+# Функция для расширения изображений (Flux Extend Image)
+def extend_image(image):
+    # Преобразуем изображение в формат, подходящий для модели
+    img = Image.open(image).convert("RGB")
+    output, _ = esrgan_model.enhance(img, outscale=4)
+    return output
+
+# Функция для генерации изображений (Image Generation)
+def generate_image(prompt):
+    with torch.autocast("cuda"):
+        image = image_generation_pipeline(prompt).images[0]
+    return image
+
 # Примеры текстов для анализа тональности
 sentiment_examples = [
     "I love programming, it's so much fun!",
@@ -305,6 +347,30 @@ with gr.Blocks() as demo:
             examples=code_fix_examples,
             examples_per_page=2
         )
+    with gr.Tab("Flux Extend Image"):
+        gr.Interface(
+            fn=extend_image,
+            inputs=gr.Image(type="pil", label="Загрузите изображение"),
+            outputs=gr.Image(type="pil", label="Расширенное изображение"),
+            title="Расширение изображений (Flux Extend Image)",
+            description="Загрузите изображение, чтобы увеличить его разрешение.",
+            examples=image_examples,
+            examples_per_page=2
+        )
+    with gr.Tab("Image Generation"):
+        gr.Interface(
+            fn=generate_image,
+            inputs=gr.Textbox(lines=2, placeholder="Введите текстовый запрос..."),
+            outputs=gr.Image(type="pil", label="Сгенерированное изображение"),
+            title="Генерация изображений (Image Generation)",
+            description="Введите текстовый запрос, чтобы сгенерировать изображение.",
+            examples=[
+                "A futuristic cityscape at night",
+                "A beautiful landscape with mountains and a lake",
+                "An astronaut riding a horse in space"
+            ],
+            examples_per_page=2
+        )
 
 # Запускаем интерфейс
 demo.launch()