Update app.py

app.py

@@ -1,52 +1,88 @@
 import gradio as gr
 import openai
 import os
-import requests
 from PIL import Image
+import requests
 from io import BytesIO
 
-# Получаем API ключ из переменных среды
-api_key = os.getenv("API_KEY")
+# Загрузим API ключ из переменных окружения
+openai.api_key = os.getenv("API_KEY")
 
-# Убедитесь, что API ключ установлен
-if api_key is None:
-    raise ValueError("API_KEY не найден. Установите переменную среды API_KEY с вашим ключом.")
-
-openai.api_key = api_key
+# Функция для генерации изображения
+def generate_image(prompt, seed, sampling, cfg, width, height):
+    response = openai.Image.create(
+        prompt=prompt, 
+        n=1, 
+        size=f"{width}x{height}",
+        seed=seed,
+        sampling=sampling,
+        cfg_scale=cfg
+    )
+    image_url = response['data'][0]['url']
+    image = Image.open(requests.get(image_url, stream=True).raw)
+    return image
 
-# Функция для генерации изображения с помощью DALL-E 3 API
-def generate_image(prompt):
-    # Создаем задачу для генерации изображения
-    response = openai.Image.create(prompt=prompt, n=1)
+# Функция для улучшения качества изображения
+def upscale_image(image, upscale_factor):
+    # Преобразуем PIL Image в Bytes, чтобы отправить его обратно в Gradio UI
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    img_str = buffered.getvalue()
+    
+    response = openai.Image.modify(
+        image=img_str,
+        n=1,
+        size=f"{image.width * upscale_factor}x{image.height * upscale_factor}"
+    )
+    upscaled_image_url = response['data'][0]['url']
+    upscaled_image = Image.open(requests.get(upscaled_image_url, stream=True).raw)
+    return upscaled_image
 
-    # Получаем файловый ID из ответа
-    file_id = response['data'][0]['id']
+# Создаем интерфейс с вкладками
+with gr.Blocks() as demo:
+    gr.Markdown("<style>footer { display: none; }</style>")  # Скрываем footer
+    with gr.Tabs():
+        with gr.TabItem("Базовые настройки"):
+            with gr.Row():
+                prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Введите описание изображения")
+        with gr.TabItem("Расширенные настройки"):
+            with gr.Row():
+                seed_input = gr.Number(label="Seed", value=0)
+                sampling_input = gr.Dropdown(label="Sampling Method", choices=["ancestral", "ddim"])
+                cfg_input = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7)
+                width_input = gr.Number(label="Width", value=256)
+                height_input = gr.Number(label="Height", value=256)
+        with gr.TabItem("Настройки апскейлинга"):
+            with gr.Row():
+                upscale_factor_input = gr.Slider(label="Upscale Factor", minimum=2, maximum=6, value=2)
 
-    # Получаем список изображений для данного файла
-    images = openai.Image.list(file_id=file_id)
+    with gr.Row():
+        generate_button = gr.Button("Генерация")
+        upscale_button = gr.Button("Улучшить качество", visible=False)
 
-    # Получаем первое изображение из списка
-    image_url = images['data'][0]['url']
+    original_image_output = gr.Image(label="Сгенерированное изображение", visible=False)
+    upscaled_image_output = gr.Image(label="Улучшенное изображение", visible=False)
 
-    # Загружаем изображение по URL и конвертируем в формат, пригодный для Gradio
-    response = requests.get(image_url)
-    image = Image.open(BytesIO(response.content))
+    # Обработчики событий кнопок
+    def generate_button_click():
+        image = generate_image(
+            prompt_input.value, 
+            seed_input.value, 
+            sampling_input.value, 
+            cfg_input.value, 
+            width_input.value, 
+            height_input.value
+        )
+        original_image_output.update(value=image, visible=True)
+        upscale_button.update(visible=True)
+        return image
 
-    return image
+    def upscale_button_click(image):
+        upscaled_image = upscale_image(image, upscale_factor_input.value)
+        upscaled_image_output.update(value=upscaled_image, visible=True)
+        return upscaled_image
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Генератор изображений DALL-E 3")
-    
-    with gr.Tabs():
-        with gr.TabItem("Базовые настройки"):
-            with gr.Column():
-                prompt_input = gr.Textbox(label="Введите в��ш запрос (prompt)", placeholder="Например: 'собака в космическом костюме'")
-                generate_button = gr.Button("Сгенерировать изображение")
-                output_image = gr.Image(label="Сгенерированное изображение")
-                generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
-        
-        with gr.TabItem("Расширенные настройки"):
-            gr.Markdown("## Скоро...")
-            # Здесь могут быть добавлены дополнительные настройки
+    generate_button.click(fn=generate_button_click, inputs=[], outputs=[original_image_output])
+    upscale_button.click(fn=upscale_button_click, inputs=[original_image_output], outputs=[upscaled_image_output])
 
 demo.launch()