import os
import tempfile
import json
from PIL import Image
import gradio as gr
import logging
import re
import time

# 패키지 임포트 방식을 변경
import google.generativeai as genai
from dotenv import load_dotenv

load_dotenv()

# Logging setup
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Gemini API 키 설정
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)
else:
    logger.warning("GEMINI_API_KEY가 설정되지 않았습니다.")

# 배경 JSON 파일 경로 설정 - 상대 경로 사용
BACKGROUNDS_DIR = "./background"

# 디렉토리 생성
if not os.path.exists(BACKGROUNDS_DIR):
    os.makedirs(BACKGROUNDS_DIR)
    logger.info(f"배경 디렉토리를 생성했습니다: {BACKGROUNDS_DIR}")
else:
    logger.info(f"배경 디렉토리가 이미 존재합니다: {BACKGROUNDS_DIR}")

# 기본 배경 데이터
SIMPLE_BACKGROUNDS = {"화이트 배경": "white background", "블랙 배경": "black background", "그라데이션 배경": "gradient background"}
STUDIO_BACKGROUNDS = {"제품 사진 스튜디오": "product photography studio", "미니멀 스튜디오": "minimal studio"}
NATURE_BACKGROUNDS = {"열대 해변": "tropical beach", "숲속": "forest", "산악 지대": "mountain landscape"}
INDOOR_BACKGROUNDS = {"모던 리빙룸": "modern living room", "세련된 사무실": "elegant office", "럭셔리 호텔": "luxury hotel"}
ABSTRACT_BACKGROUNDS = {"네온 조명": "neon lights", "추상적 그라데이션": "abstract gradient", "우주 배경": "space background"}

# JSON 파일 로드 함수
def load_background_json(filename):
    file_path = os.path.join(BACKGROUNDS_DIR, filename)
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            logger.info(f"{filename} 파일을 성공적으로 로드했습니다. {len(data)} 항목 포함.")
            return data
    except FileNotFoundError:
        logger.warning(f"경고: {filename} 파일을 찾을 수 없습니다. 기본값을 사용합니다.")
        # 기본값 생성 - 파일이 없는 경우 생성
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                if "simple" in filename:
                    json.dump(SIMPLE_BACKGROUNDS, f, ensure_ascii=False, indent=2)
                elif "studio" in filename:
                    json.dump(STUDIO_BACKGROUNDS, f, ensure_ascii=False, indent=2)
                elif "nature" in filename:
                    json.dump(NATURE_BACKGROUNDS, f, ensure_ascii=False, indent=2)
                elif "indoor" in filename:
                    json.dump(INDOOR_BACKGROUNDS, f, ensure_ascii=False, indent=2)
                elif "abstract" in filename:
                    json.dump(ABSTRACT_BACKGROUNDS, f, ensure_ascii=False, indent=2)
            logger.info(f"{filename} 파일을 기본값으로 생성했습니다.")
        except Exception as e:
            logger.error(f"파일 생성 중 오류: {str(e)}")
        
        # 기본값 반환
        if "simple" in filename:
            return SIMPLE_BACKGROUNDS
        elif "studio" in filename:
            return STUDIO_BACKGROUNDS
        elif "nature" in filename:
            return NATURE_BACKGROUNDS
        elif "indoor" in filename:
            return INDOOR_BACKGROUNDS
        elif "abstract" in filename:
            return ABSTRACT_BACKGROUNDS
        return {}
    except json.JSONDecodeError:
        logger.warning(f"경고: {filename} 파일의 JSON 형식이 올바르지 않습니다. 기본값을 사용합니다.")
        return {}
    except Exception as e:
        logger.warning(f"경고: {filename} 파일 로드 중 오류 발생: {str(e)}. 기본값을 사용합니다.")
        return {}

# 배경 데이터 로드
SIMPLE_BACKGROUNDS = load_background_json("simple_backgrounds.json")
STUDIO_BACKGROUNDS = load_background_json("studio_backgrounds.json")
NATURE_BACKGROUNDS = load_background_json("nature_backgrounds.json")
INDOOR_BACKGROUNDS = load_background_json("indoor_backgrounds.json")
ABSTRACT_BACKGROUNDS = load_background_json("abstract_backgrounds.json")

def save_binary_file(file_name, data):
    with open(file_name, "wb") as f:
        f.write(data)

def translate_prompt_to_english(prompt):
    if not re.search("[가-힣]", prompt):
        return prompt
    
    prompt = prompt.replace("#1", "IMAGE_TAG_ONE")
    
    try:
        if not GEMINI_API_KEY:
            logger.error("Gemini API 키가 설정되지 않았습니다.")
            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
            return prompt
            
        translation_prompt = f"""
        Translate the following Korean text to English:
        
        {prompt}
        
        IMPORTANT: The token IMAGE_TAG_ONE is a special tag 
        and must be preserved exactly as is in your translation. Do not translate this token.
        """
        
        logger.info(f"Translation prompt: {translation_prompt}")
        
        # 새로운 API 사용 방식
        response = genai.generate_text(
            model="gemini-2.0-flash",
            prompt=translation_prompt,
            temperature=0.2,
            top_p=0.95,
            top_k=40,
            max_output_tokens=512
        )
        
        translated_text = response.result
        
        if translated_text and translated_text.strip():
            translated_text = translated_text.replace("IMAGE_TAG_ONE", "#1")
            logger.info(f"Translated text: {translated_text.strip()}")
            return translated_text.strip()
        else:
            logger.warning("번역 결과가 없습니다. 원본 프롬프트 사용")
            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
            return prompt
    except Exception as e:
        logger.exception("번역 중 오류 발생:")
        prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
        return prompt

def preprocess_prompt(prompt, image1):
    has_img1 = image1 is not None

    if "#1" in prompt and not has_img1:
        prompt = prompt.replace("#1", "첫 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#1", "첫 번째 이미지")

    prompt += " 이미지를 생성해주세요. 이미지에 텍스트나 글자를 포함하지 마세요."
    return prompt

def generate_with_images(prompt, images, variation_index=0):
    try:
        if not GEMINI_API_KEY:
            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."

        client = genai.Client(api_key=GEMINI_API_KEY)
        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}, 변형 인덱스: {variation_index}")

        variation_suffixes = [
            " Create this as the first variation. Do not add any text, watermarks, or labels to the image.",
            " Create this as the second variation with more vivid colors. Do not add any text, watermarks, or labels to the image.",
            " Create this as the third variation with a more creative style. Do not add any text, watermarks, or labels to the image.",
            " Create this as the fourth variation with enhanced details. Do not add any text, watermarks, or labels to the image."
        ]
        
        if variation_index < len(variation_suffixes):
            prompt = prompt + variation_suffixes[variation_index]
        else:
            prompt = prompt + " Do not add any text, watermarks, or labels to the image."

        contents = [prompt]
        for idx, img in enumerate(images, 1):
            if img is not None:
                contents.append(img)
                logger.info(f"이미지 #{idx} 추가됨")

        response = client.models.generate_content(
            model="gemini-2.0-flash-exp-image-generation",
            contents=contents,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image'],
                temperature=1,
                top_p=0.95,
                top_k=40,
                max_output_tokens=8192
            )
        )

        # 임시 파일은 항상 JPG 확장자로 생성
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            temp_path = tmp.name
            result_text = ""
            image_found = False
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    result_text += part.text
                    logger.info(f"응답 텍스트: {part.text}")
                elif hasattr(part, 'inline_data') and part.inline_data:
                    save_binary_file(temp_path, part.inline_data.data)
                    image_found = True
                    logger.info("응답에서 이미지 추출 성공")
            if not image_found:
                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
            
            result_img = Image.open(temp_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")  # JPG는 투명도를 지원하지 않으므로 RGB로 변환
            
            # 변환된 이미지를 JPG로 저장
            result_img.save(temp_path, format="JPEG", quality=95)
            
            # 파일 경로 반환
            return temp_path, f"이미지가 성공적으로 생성되었습니다. {result_text}"
    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"

def process_images_with_prompt(image1, prompt, variation_index=0, max_retries=3):
    retry_count = 0
    last_error = None
    
    while retry_count < max_retries:
        try:
            images = [image1]
            valid_images = [img for img in images if img is not None]
            if not valid_images:
                return None, "이미지를 업로드해주세요.", ""

            if prompt and prompt.strip():
                processed_prompt = preprocess_prompt(prompt, image1)
                if re.search("[가-힣]", processed_prompt):
                    final_prompt = translate_prompt_to_english(processed_prompt)
                else:
                    final_prompt = processed_prompt
            else:
                final_prompt = "Please creatively transform this image into a more vivid and artistic version. Do not include any text or watermarks in the generated image."
                logger.info("Default prompt generated for single image")

            result_img, status = generate_with_images(final_prompt, valid_images, variation_index)
            if result_img is not None:
                return result_img, status, final_prompt
            else:
                last_error = status
                retry_count += 1
                logger.warning(f"이미지 생성 실패, 재시도 {retry_count}/{max_retries}: {status}")
                time.sleep(1)
        except Exception as e:
            last_error = str(e)
            retry_count += 1
            logger.exception(f"이미지 처리 중 오류 발생, 재시도 {retry_count}/{max_retries}:")
            time.sleep(1)
    
    return None, f"최대 재시도 횟수({max_retries}회) 초과 후 실패: {last_error}", prompt

def generate_multiple_images(image1, prompt, progress=gr.Progress()):
    results = []
    statuses = []
    prompts = []
    
    num_images = 4
    max_retries = 3
    
    progress(0, desc="이미지 생성 준비 중...")
    
    for i in range(num_images):
        progress((i / num_images), desc=f"{i+1}/{num_images} 이미지 생성 중...")
        result_img, status, final_prompt = process_images_with_prompt(image1, prompt, i, max_retries)
        
        if result_img is not None:
            results.append(result_img)
            statuses.append(f"이미지 #{i+1}: {status}")
            prompts.append(f"이미지 #{i+1}: {final_prompt}")
        else:
            results.append(None)
            statuses.append(f"이미지 #{i+1} 생성 실패: {status}")
            prompts.append(f"이미지 #{i+1}: {final_prompt}")
        
        time.sleep(1)
    
    progress(1.0, desc="이미지 생성 완료!")
    
    while len(results) < 4:
        results.append(None)
    
    combined_status = "\n".join(statuses)
    combined_prompts = "\n".join(prompts)
    
    return results[0], results[1], results[2], results[3], combined_status, combined_prompts

def generate_system_instruction():
    return """당신은 상품 이미지의 배경을 변경하기 위한 고품질 프롬프트를 생성하는 전문가입니다.
사용자가 제공하는 상품명, 배경 유형, 추가 요청사항을 바탕으로 미드저니(Midjourney)에 사용할 수 있는 
상세하고 전문적인 프롬프트를 영어로 생성해주세요.
다음 가이드라인을 반드시 따라야 합니다:
1. 상품을 "#1"로 지정하여 참조합니다. (예: "skincare tube (#1)")
2. *** 매우 중요: 상품의 원래 특성(디자인, 색상, 형태, 로고, 패키지 등)은 어떤 상황에서도 절대 변경하지 않습니다. ***
3. *** 상품의 본질적 특성은 유지하되, 자연스러운 환경 통합을 위한 조명과 그림자는 허용합니다: ***
   - 상품 자체의 색상, 디자인, 형태, 텍스처는 절대 수정하지 않습니다.
   - 환경과 자연스럽게 어울리는 그림자, 주변 조명 효과는 허용됩니다.
   - 상품에 물방울, 응축, 금, 은과 같은 추가 요소나 물리적 효과는 적용하지 않습니다.
   - 환경에 어울리는 자연스러운 빛 반사, 주변 조명, 그림자는 사실적 통합감을 위해 적용할 수 있습니다.
4. 이미지 비율은 정확히 1:1(정사각형) 형식으로 지정합니다. 프롬프트에 "square format", "1:1 ratio" 또는 "aspect ratio 1:1"을 명시적으로 포함합니다.
5. 상품은 반드시 정사각형 구도의 정중앙에 배치되어야 합니다.
6. 상품을 이미지의 주요 초점으로 부각시키고, 상품의 비율이 전체 이미지에서 크게 차지하도록 합니다.
7. 상품 이미지 컷아웃(#1)의 기본 형태와 색상은 유지하면서, 선택한 환경에 자연스럽게 통합되도록 합니다.
8. 고급스러운 상업적 이미지를 위한 다음 환경 요소들을 포함하세요:
   - 상품과 어울리는 주변 환경/배경 요소를 추가합니다. 예를 들어, 화장품 주변에 꽃이나 허브, 음료 제품 옆에 과일, 전자제품 근처에 현대적 소품 등.
   - 환경의 조명 효과(림 라이트, 백라이트, 소프트박스 등)를 설명합니다.
   - 상품이 환경에 자연스럽게 존재하는 것처럼 보이도록 적절한 그림자와 빛 표현을 포함합니다.
   - 상품의 용도나 장점을 간접적으로 암시하는 배경 요소를 포함합니다.
   - 프로페셔널한 상업 사진 효과(선택적 피사계 심도, 소프트 포커스, 스튜디오 조명 등)를 명시합니다.
9. 프롬프트에 다음 요소들을 명시적으로 포함하세요:
   - "highly detailed commercial photography"
   - "award-winning product photography"
   - "professional advertising imagery"
   - "studio quality"
   - "magazine advertisement quality"
10. 배경 환경 요소를 상품 카테고리에 맞게 선택합니다:
    - 스킨케어 제품: 깨끗한 욕실 선반, 우아한 화장대, 스파 같은 환경 등
    - 음료 제품: 세련된 테이블, 파티 환경, 야외 피크닉 장면 등
    - 전자 제품: 세련된 작업 공간, 현대적인 거실, 미니멀한 책상 등
    - 패션/의류: 세련된 쇼룸, 도시 거리, 엘레강스한 라이프스타일 환경 등
    - 식품 제품: 깔끔한 주방, 식탁, 요리 환경 등
11. 사용자가 제공한 구체적인 배경과 추가 요청사항을 정확히 반영합니다.
12. 프롬프트는 미드저니 AI에 최적화되어야 합니다.
13. 프롬프트 끝에 "--ar 1:1 --s 750 --q 2" 파라미터를 추가하여 미드저니에서 고품질 정사각형 비율을 강제합니다.
출력 형식은 영어로 된 단일 단락의 상세한 프롬프트여야 하며, 끝에 미드저니 파라미터가 포함되어야 합니다.
"""

def generate_prompt_with_gemini(product_name, background_info, additional_info=""):
    if not GEMINI_API_KEY:
        return "Gemini API 키가 설정되지 않았습니다. 환경 변수 GEMINI_API_KEY를 설정하거나 코드에 직접 입력하세요."
    try:
        prompt_request = f"""
상품명: {product_name}
배경 유형: {background_info.get('english', 'studio')}
배경 카테고리: {background_info.get('category', '')}
배경 이름: {background_info.get('name', '')}
추가 요청사항: {additional_info}
중요 요구사항:
1. 상품이 크게 부각되고 이미지에서 중심적인 위치를 차지하도록 프롬프트를 생성해주세요.
2. 이미지는 정확히 1:1 비율(정사각형)이어야 합니다.
3. 상품은 정사각형 프레임의 정중앙에 위치해야 합니다.
4. 상품의 디자인, 색상, 형태, 로고 등 본질적 특성은 절대 수정하지 마세요.
5. 환경과의 자연스러운 통합을 위한 조명 효과와 그림자는 포함해주세요.
6. 상품을 더 돋보이게 하는 배경 환경을 설명해주세요.
7. 고급스러운 상업 광고 품질의 이미지가 되도록 환경 설명을 해주세요.
8. 프롬프트 끝에 미드저니 파라미터 "--ar 1:1 --s 750 --q 2"를 추가해주세요.
한국어 입력 내용을 영어로 적절히 번역하여 반영해주세요.
"""
        
        # google-generativeai 패키지 사용 방식으로 변경
        system_instruction = generate_system_instruction()
        
        # GenerativeModel 대신 genai.chat 사용
        chat = genai.chat(
            model="gemini-2.0-flash",
            messages=[
                {"role": "system", "content": system_instruction},
                {"role": "user", "content": prompt_request}
            ],
            generation_config={
                "temperature": 0.7,
                "top_p": 0.95,
                "top_k": 64,
                "max_output_tokens": 1024,
            }
        )
        
        response_text = chat.last.text.strip()
        
        if "--ar 1:1" not in response_text:
            response_text = response_text.rstrip(".") + ". --ar 1:1 --s 750 --q 2"
            
        # 미드저니 파라미터 제거 (Gemini API와는 달리 불필요하므로)
        response_text = response_text.replace(" --ar 1:1 --s 750 --q 2", "")
        
        return response_text
    except Exception as e:
        return f"프롬프트 생성 중 오류가 발생했습니다: {str(e)}"

# 선택된 배경 정보 가져오기
def get_selected_background_info(bg_type, simple, studio, nature, indoor, abstract):
    if bg_type == "심플 배경":
        return {
            "category": "심플 배경",
            "name": simple,
            "english": SIMPLE_BACKGROUNDS.get(simple, "white background")
        }
    elif bg_type == "스튜디오 배경":
        return {
            "category": "스튜디오 배경",
            "name": studio,
            "english": STUDIO_BACKGROUNDS.get(studio, "product photography studio")
        }
    elif bg_type == "자연 환경":
        return {
            "category": "자연 환경",
            "name": nature,
            "english": NATURE_BACKGROUNDS.get(nature, "natural environment")
        }
    elif bg_type == "실내 환경":
        return {
            "category": "실내 환경",
            "name": indoor,
            "english": INDOOR_BACKGROUNDS.get(indoor, "indoor environment")
        }
    elif bg_type == "추상/특수 배경":
        return {
            "category": "추상/특수 배경",
            "name": abstract,
            "english": ABSTRACT_BACKGROUNDS.get(abstract, "abstract background")
        }
    else:
        return {
            "category": "기본 배경",
            "name": "화이트 배경",
            "english": "white background"
        }

# 프롬프트 생성 함수
def create_prompt(image, bg_type, simple, studio, nature, indoor, abstract, product_text, additional_text):
    if image is None:
        return "이미지를 업로드해주세요."
    
    product_text = product_text.strip() or "제품"
    
    # 배경 정보 가져오기
    background_info = get_selected_background_info(bg_type, simple, studio, nature, indoor, abstract)
    
    try:
        prompt = generate_prompt_with_gemini(product_text, background_info, additional_text)
        return prompt
    except Exception as e:
        error_msg = f"프롬프트 생성 중 오류 발생: {str(e)}"
        return error_msg


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            # 이미지 업로드 섹션
            gr.Markdown("## 이미지 업로드")
            with gr.Row():
                image1_input = gr.Image(type="pil", label="이미지 업로드", image_mode="RGB")
            
            # 제품명 입력 섹션
            gr.Markdown("## 제품 정보")
            product_name = gr.Textbox(label="제품명 (한국어 입력)", placeholder="예: 스킨케어 튜브, 텀블러 등")
            
            # 배경 유형 섹션
            gr.Markdown("## 배경 설정")
            background_type = gr.Radio(
                choices=["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "추상/특수 배경"],
                label="배경 유형",
                value="심플 배경"
            )
            
            # 각 배경 유형에 맞는 드롭다운 컴포넌트들
            simple_dropdown = gr.Dropdown(
                choices=list(SIMPLE_BACKGROUNDS.keys()),
                value=list(SIMPLE_BACKGROUNDS.keys())[0] if SIMPLE_BACKGROUNDS else None,
                label="심플 배경 선택",
                visible=True,
                interactive=True
            )
            
            studio_dropdown = gr.Dropdown(
                choices=list(STUDIO_BACKGROUNDS.keys()),
                value=list(STUDIO_BACKGROUNDS.keys())[0] if STUDIO_BACKGROUNDS else None,
                label="스튜디오 배경 선택",
                visible=False,
                interactive=True
            )
            
            nature_dropdown = gr.Dropdown(
                choices=list(NATURE_BACKGROUNDS.keys()),
                value=list(NATURE_BACKGROUNDS.keys())[0] if NATURE_BACKGROUNDS else None,
                label="자연 환경 선택",
                visible=False,
                interactive=True
            )
            
            indoor_dropdown = gr.Dropdown(
                choices=list(INDOOR_BACKGROUNDS.keys()),
                value=list(INDOOR_BACKGROUNDS.keys())[0] if INDOOR_BACKGROUNDS else None,
                label="실내 환경 선택",
                visible=False,
                interactive=True
            )
            
            abstract_dropdown = gr.Dropdown(
                choices=list(ABSTRACT_BACKGROUNDS.keys()),
                value=list(ABSTRACT_BACKGROUNDS.keys())[0] if ABSTRACT_BACKGROUNDS else None,
                label="추상/특수 배경 선택",
                visible=False,
                interactive=True
            )
            
            # 추가 요청사항
            additional_info = gr.Textbox(
                label="추가 요청사항 (선택사항)",
                placeholder="예: 고급스러운 느낌, 밝은 조명, 자연스러운 보조적인 객체를 추가해주세요 등",
                lines=2
            )
            
            # 프롬프트 입력 & 생성 섹션
            gr.Markdown("## 프롬프트 설정")
            
            generate_prompt_btn = gr.Button("프롬프트 생성")
            
            prompt_input = gr.Textbox(
                lines=3,
                placeholder="여기에 프롬프트가 생성됩니다. 수동으로 수정할 수도 있습니다. '#1'으로 업로드한 이미지를 참조합니다.",
                label="생성된 프롬프트"
            )
            
            with gr.Row():
                submit_single_btn = gr.Button('이미지 생성 (1장)')
                submit_btn = gr.Button('이미지 생성 (4장)')
        
        with gr.Column():
            gr.Markdown("## 생성된 이미지")
            with gr.Row():
                with gr.Column():
                    output_image1 = gr.Image(label="이미지 #1", type="filepath")
                    output_image3 = gr.Image(label="이미지 #3", type="filepath")
                with gr.Column():
                    output_image2 = gr.Image(label="이미지 #2", type="filepath")
                    output_image4 = gr.Image(label="이미지 #4", type="filepath")
            
            output_text = gr.Textbox(label="결과 정보", lines=2)
            prompt_display = gr.Textbox(label="사용된 프롬프트 (영어)", lines=2)

    # 배경 유형에 따라 드롭다운 표시 업데이트
    def update_dropdowns(bg_type):
        return {
            simple_dropdown: gr.update(visible=(bg_type == "심플 배경")),
            studio_dropdown: gr.update(visible=(bg_type == "스튜디오 배경")),
            nature_dropdown: gr.update(visible=(bg_type == "자연 환경")),
            indoor_dropdown: gr.update(visible=(bg_type == "실내 환경")),
            abstract_dropdown: gr.update(visible=(bg_type == "추상/특수 배경"))
        }

    # 배경 유형 변경 이벤트 처리
    background_type.change(
        fn=update_dropdowns,
        inputs=[background_type],
        outputs=[simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, abstract_dropdown]
    )

    # 프롬프트 생성 버튼 클릭 이벤트
    generate_prompt_btn.click(
        fn=create_prompt,
        inputs=[
            image1_input,
            background_type,
            simple_dropdown,
            studio_dropdown,
            nature_dropdown,
            indoor_dropdown,
            abstract_dropdown,
            product_name,
            additional_info
        ],
        outputs=[prompt_input]
    )
    
    # 단일 이미지 생성 버튼 이벤트 연결
    submit_single_btn.click(
        fn=lambda image1, prompt: process_images_with_prompt(image1, prompt, 0),
        inputs=[image1_input, prompt_input],
        outputs=[output_image1, output_text, prompt_display],
    )
    
    # 4장 이미지 생성 버튼 이벤트 연결
    submit_btn.click(
        fn=generate_multiple_images,
        inputs=[image1_input, prompt_input],
        outputs=[output_image1, output_image2, output_image3, output_image4, output_text, prompt_display],
    )

demo.queue()
demo.launch()