Update app.py

app.py

@@ -263,15 +263,36 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
             error_message = "Model is currently loading. Please try again in a few moments."
         raise gr.Error(error_message)
 
-def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height):
+def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height, progress=gr.Progress()):
     if len(selected_models) > 4:
         raise gr.Error("Please select up to 4 models")
     if len(selected_models) == 0:
         raise gr.Error("Please select at least 1 model")
     
-    # 초기 이미지 배열 생성 (None으로 채움)
+    # 초기 이미지 배열 생성
     images = [None] * 4
     
+    # 각 모델별로 이미지 생성
+    for idx, model_name in enumerate(selected_models[:4]):
+        try:
+            progress(idx / len(selected_models), f"Generating image for {model_name}...")
+            img = query(prompt, model_name, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height)
+            images[idx] = img
+        except Exception as e:
+            print(f"Error generating image for {model_name}: {str(e)}")
+            continue
+    
+    # 남은 슬롯을 마지막 생성된 이미지로 채움
+    last_valid_image = next((img for img in reversed(images) if img is not None), None)
+    if last_valid_image:
+        for i in range(len(images)):
+            if images[i] is None:
+                images[i] = last_valid_image
+    
+    progress(1.0, "Generation complete!")
+    return images
+
+    
     def update_gallery():
         return images
     
@@ -409,8 +430,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
             height="auto"
         )
 
-
-
     generate_btn.click(
         fn=generate_grid,
         inputs=[
@@ -425,10 +444,10 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
             width,
             height
         ],
-        outputs=gallery,
-        every=True  # 중간 결과를 실시간으로 업데이트
+        outputs=gallery
     )
 
+
     
     def filter_models(search_term):
         filtered_models = [m for m in models_list if search_term.lower() in m.lower()]