Update app.py

app.py

@@ -13,6 +13,9 @@ model_repo_id = "stabilityai/sdxl-turbo"
 
 torch_dtype = torch.float32  # Тип данных для CPU
 
+
+
+# Список доступных моделей
 models = [
     'stabilityai/stable-diffusion-3.5-large-turbo',
     'black-forest-labs/FLUX.1-dev',
@@ -42,170 +45,99 @@ models = [
     'Jonjew/NSFWMaster',
 ]
 
-
-
-# Функция загрузки моделей
+# Загрузка моделей
 def load_fn(models):
     global models_load
     models_load = {}
     for model in models:
-        if model not in models_load.keys():
-            try:
-                m = gr.load(f'models/{model}')
-            except Exception as error:
-                m = gr.Interface(lambda txt: None, ['text'], ['image'])
-            models_load.update({model: m})
-
-# Инициализация моделей
-models = ["Model1", "Model2", "Model3", "Model4"]  # Список доступных моделей
+        try:
+            models_load[model] = gr.load(f'models/{model}')
+        except Exception as e:
+            print(f"Failed to load model {model}: {e}")
+            models_load[model] = gr.Interface(lambda txt: None, ['text'], ['image'])
+
 load_fn(models)
 
 num_models = len(models)
-default_models = models[:num_models]
+default_models = models[:4]  # Выбираем первые 4 модели по умолчанию
 
-# Функция расширения списка выборов
+# Расширение списка для выбора
 def extend_choices(choices):
     return choices + (num_models - len(choices)) * ['NA']
 
-# Обновление изображений
+# Обновление изображения
 def update_imgbox(choices, width, height):
     choices_plus = extend_choices(choices)
     return [
         gr.Image(
-            None, 
-            label=m, 
-            visible=(m != 'NA'), 
-            elem_id="custom_image", 
-            width=width, 
+            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, width, height, max_retries=3):
-    """
-    Генерация изображения с обработкой ошибок и повторными попытками.
-    """
-    if model_str == 'NA':
+# Функция генерации
+def gen_fn(model_str, prompt, width, height):
+    if model_str == 'NA' or model_str not in models_load:
+        return None
+    try:
+        result = models_load[model_str](prompt)
+        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 for model {model_str}: {e}")
         return None
-    
-    noise = str(randint(0, 9999999))
-    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:',
                 lines=3,
                 container=False,
                 elem_id="custom_textbox",
-                placeholder="Prompt"
+                placeholder="Enter your prompt here"
             )
-            # Выбор ширины и высоты
             width_slider = gr.Slider(
-                minimum=100, 
-                maximum=1000, 
-                step=50, 
-                value=250, 
-                label="Width"
+                minimum=100, maximum=1000, step=50, value=250, label="Width"
             )
             height_slider = gr.Slider(
-                minimum=100, 
-                maximum=1000, 
-                step=50, 
-                value=250, 
-                label="Height"
+                minimum=100, maximum=1000, step=50, value=250, label="Height"
             )
-            with gr.Row(): 
+            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")
-                
-                def on_generate_click():
-                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
-                
-                def on_stop_click():
-                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
-                
-                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, 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, 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, width_slider, height_slider], o)
-            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
+        # Привязка генерации к кнопке
+        for i, model in enumerate(default_models):
+            row = output_row1 if i < half else output_row2
+            row[i % half] = gr.Image(label=model, elem_id="custom_image")
+            gen_button.click(gen_fn, inputs=[model, txt_input, width_slider, height_slider], outputs=row[i % half])
     
-    # Чекбоксы выбора моделей в аккордеоне
     with gr.Accordion('Model selection', elem_id="custom_accordion"):
-        half = len(models) // 2
-        model_choice_row1 = models[:half]
-        model_choice_row2 = models[half:]
-        
-        model_choice1 = gr.CheckboxGroup(
-            model_choice_row1,
-            label=f'{len(model_choice_row1)} models selected',
-            value=default_models[:half],
-            interactive=True,
+        model_choice = gr.CheckboxGroup(
+            choices=models,
+            label=f'{len(models)} available models',
+            value=default_models,
             elem_id="custom_checkbox_group"
         )
-        model_choice2 = gr.CheckboxGroup(
-            model_choice_row2,
-            label=f'{len(model_choice_row2)} models selected',
-            value=default_models[half:],
-            interactive=True,
-            elem_id="custom_checkbox_group"
+        model_choice.change(
+            update_imgbox,
+            inputs=[model_choice, width_slider, height_slider],
+            outputs=[output_row1, output_row2]
         )
-        
-        # Обновление первого ряда
-        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, width_slider, height_slider], output_row2)
-        model_choice2.change(extend_choices, model_choice2, current_models_row2)
-
-
 
     with gr.Row():
         gr.HTML("")