Update app.py

app.py

@@ -1,3 +1,4 @@
+from PIL import Image
 import gradio as gr
 from random import randint
 
@@ -41,6 +42,9 @@ models = [
     'Jonjew/NSFWMaster',
 ]
 
+
+
+# Функция загрузки моделей
 def load_fn(models):
     global models_load
     models_load = {}
@@ -52,26 +56,72 @@ def load_fn(models):
                 m = gr.Interface(lambda txt: None, ['text'], ['image'])
             models_load.update({model: m})
 
+# Инициализация моделей
+models = ["Model1", "Model2", "Model3", "Model4"]  # Список доступных моделей
 load_fn(models)
 
 num_models = len(models)
 default_models = models[:num_models]
 
+# Функция расширения списка выборов
 def extend_choices(choices):
     return choices + (num_models - len(choices)) * ['NA']
 
-def update_imgbox(choices):
+# Обновление изображений
+def update_imgbox(choices, width, height):
     choices_plus = extend_choices(choices)
-    return [gr.Image(None, label=m, visible=(m != 'NA'), elem_id="custom_image") for m in choices_plus]
+    return [
+        gr.Image(
+            None, 
+            label=m, 
+            visible=(m != 'NA'), 
+            elem_id="custom_image", 
+            width=width, 
+            height=height
+        ) for m in choices_plus
+    ]
 
-def gen_fn(model_str, prompt):
+# Генерация изображений
+def gen_fn(model_str, prompt, width, height, max_retries=3):
+    """
+    Генерация изображения с обработкой ошибок и повторными попытками.
+    """
     if model_str == 'NA':
         return None
+    
     noise = str(randint(0, 9999999))
-    return models_load[model_str](f'{prompt} {noise}')
+    retries = 0
+
+    while retries < max_retries:
+        try:
+            result = models_load[model_str](f'{prompt} {noise}')
+            
+            # Проверка на тип tuple
+            if isinstance(result, tuple):
+                print(f"Result is a tuple, retrying... Attempt {retries + 1}")
+                retries += 1
+                continue
+            
+            # Если результат строка (путь к файлу), открываем изображение
+            if isinstance(result, str):
+                result = Image.open(result)
+            
+            # Изменение размера изображения
+            if isinstance(result, Image.Image):
+                result = result.resize((width, height))
+            
+            return result
+        except Exception as e:
+            print(f"Error generating image: {e}")
+            retries += 1
+
+    print("Failed to generate image after multiple retries.")
+    return None
 
+# Создание интерфейса
 def make_me():
     with gr.Row():
+        # Ввод текста
         with gr.Column(scale=1):
             txt_input = gr.Textbox(
                 label='Your prompt:',
@@ -80,6 +130,21 @@ def make_me():
                 elem_id="custom_textbox",
                 placeholder="Prompt"
             )
+            # Выбор ширины и высоты
+            width_slider = gr.Slider(
+                minimum=100, 
+                maximum=1000, 
+                step=50, 
+                value=250, 
+                label="Width"
+            )
+            height_slider = gr.Slider(
+                minimum=100, 
+                maximum=1000, 
+                step=50, 
+                value=250, 
+                label="Height"
+            )
             with gr.Row(): 
                 gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
                 stop_button = gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
@@ -93,25 +158,25 @@ def make_me():
                 gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
                 stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
     
-    # Делим вывод изображений на два ряда
+    # Вывод изображений в два ряда
     with gr.Row():
         half = len(default_models) // 2
-        output_row1 = [gr.Image(label=m, min_width=250, height=250, elem_id="custom_image") for m in default_models[:half]]
-        output_row2 = [gr.Image(label=m, min_width=250, height=250, elem_id="custom_image") for m in default_models[half:]]
+        output_row1 = [gr.Image(label=m, elem_id="custom_image") for m in default_models[:half]]
+        output_row2 = [gr.Image(label=m, elem_id="custom_image") for m in default_models[half:]]
         current_models_row1 = [gr.Textbox(m, visible=False) for m in default_models[:half]]
         current_models_row2 = [gr.Textbox(m, visible=False) for m in default_models[half:]]
         
-        # Первый ряд изображений
+        # Генерация для первого ряда
         for m, o in zip(current_models_row1, output_row1):
-            gen_event = gen_button.click(gen_fn, [m, txt_input], o)
+            gen_event = gen_button.click(gen_fn, [m, txt_input, width_slider, height_slider], o)
             stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
         
-        # Второй ряд изображений
+        # Генерация для второго ряда
         for m, o in zip(current_models_row2, output_row2):
-            gen_event = gen_button.click(gen_fn, [m, txt_input], o)
+            gen_event = gen_button.click(gen_fn, [m, txt_input, width_slider, height_slider], o)
             stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
     
-    # Разделяем чекбоксы на два ряда
+    # Чекбоксы выбора моделей в аккордеоне
     with gr.Accordion('Model selection', elem_id="custom_accordion"):
         half = len(models) // 2
         model_choice_row1 = models[:half]
@@ -132,14 +197,16 @@ def make_me():
             elem_id="custom_checkbox_group"
         )
         
-        # Обновляем первый ряд
-        model_choice1.change(update_imgbox, model_choice1, output_row1)
+        # Обновление первого ряда
+        model_choice1.change(update_imgbox, [model_choice1, width_slider, height_slider], output_row1)
         model_choice1.change(extend_choices, model_choice1, current_models_row1)
         
-        # Обновляем второй ряд
-        model_choice2.change(update_imgbox, model_choice2, output_row2)
+        # Обновление второго ряда
+        model_choice2.change(update_imgbox, [model_choice2, width_slider, height_slider], output_row2)
         model_choice2.change(extend_choices, model_choice2, current_models_row2)
 
+
+
     with gr.Row():
         gr.HTML("")