Update app.py

app.py

@@ -5,10 +5,6 @@ import random
 import os
 from PIL import Image
 import json
-from dotenv import load_dotenv
-
-# Load environment variables
-load_dotenv()
 
 # Get API token from environment variable 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -267,14 +263,33 @@ 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):
+    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")
+    
+    images = []
+    for model_name in selected_models[:4]:  # 최대 4개 모델까지만 처리
+        try:
+            img = query(prompt, model_name, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height)
+            images.append(img)
+        except Exception as e:
+            print(f"Error generating image for {model_name}: {str(e)}")
+            continue
+    
+    # 4개 이미지가 되도록 부족한 만큼 마지막 이미지를 복제
+    while len(images) < 4:
+        images.append(images[-1] if images else None)
+    
+    return images
+
 css = """
 footer {
     visibility: hidden;
 }
 """
 
-print("Initializing Gradio interface...")
-
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
     gr.Markdown("# AI Image Generator")
     
@@ -319,16 +334,17 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     lines=1
                 )
                 
-                # 상위 9개 모델을 기본으로 설정
+                # 상위 4개 모델을 기본으로 설정
                 default_models = [
+                    "Seawolf Teeeee",
                     "Stable Diffusion 3.5 Large",
                     "Stable Diffusion 3.5 Large Turbo",
-                    "Stable Diffusion XL",
-                    "FLUX.1 [Schnell]"
+                    "Stable Diffusion XL"
                 ]
 
                 # 전체 모델 리스트
                 models_list = [
+                    "Seawolf Teeeee",
                     "Stable Diffusion 3.5 Large",
                     "Stable Diffusion 3.5 Large Turbo",
                     "Stable Diffusion XL",
@@ -363,10 +379,10 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     "Collage Flux"
                 ]
 
-                model = gr.Checkboxgroup(  # Radio를 Checkboxgroup으로 변경
+                model = gr.Checkboxgroup(
                     label="Select Models (Choose up to 4)",
                     choices=models_list,
-                    value=default_models,  # 기본값으로 상위 4개 선택
+                    value=default_models,
                     interactive=True
                 )