app.py

"""
PixelFree - AI Image Processing Service
HuggingFace Spaces용 Gradio 애플리케이션
Polygom API를 활용한 이미지 처리 서비스

Features:
- 배경 제거 (Background Removal)
- 배경 생성 (Background Generation) 
- 얼굴 교체 (Face Swap)

Version: 3.0.0
Version 3: Simplified Multi-Mode Workflow
Author: PixelFree Team
"""

import gradio as gr
import requests
import base64
import io
import logging
import tempfile
from PIL import Image
from typing import Optional, Tuple
from datetime import datetime
from collections import OrderedDict
import time
import os

# ===== Configuration =====
API_BASE_URL = os.environ.get("API_BASE_URL")
API_ENDPOINTS = {
    "background_removal": os.environ.get("BACKGROUND_REMOVAL_ENDPOINT"),
    "background_generation": os.environ.get("BACKGROUND_GENERATION_ENDPOINT"),
    "face_swap": os.environ.get("FACE_SWAP_ENDPOINT"),
    "video_generation": os.environ.get("VIDEO_GEN"),
}
VIDEO_PROMPT = os.environ.get("VIDEO_PROMPT", "환하게 웃으며 눈을 깜빡인다.")

# 허용된 사용자 목록 (Hugging Face Secrets에서 가져오기)
ALLOWED_USERS = os.environ.get("ALLOWED_USERS", "").split(",")
ALLOWED_USERS = [user.strip() for user in ALLOWED_USERS if user.strip()]

IMAGE_HEIGHT = 500
GALLERY_COLUMNS = 3
GALLERY_ROWS = 3
logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)

IMAGE_CACHE_LIMIT = 24
image_cache: "OrderedDict[str, Image.Image]" = OrderedDict()


def cache_image(path: str, image: Image.Image) -> None:
    """Add image to global cache with simple LRU eviction."""
    if path in image_cache:
        image_cache.move_to_end(path)
    image_cache[path] = image
    while len(image_cache) > IMAGE_CACHE_LIMIT:
        image_cache.popitem(last=False)

output_history = []


def ensure_png_image(image: Optional[Image.Image]) -> Optional[Image.Image]:
    """Ensure provided PIL image is stored as a PNG image (returns copy)."""
    if image is None:
        return None

    working_image = image
    if working_image.mode not in ("RGB", "RGBA"):
        working_image = working_image.convert("RGBA")

    buffer = io.BytesIO()
    # Convert to PNG while preserving alpha when present
    target_mode = "RGBA" if working_image.mode == "RGBA" else "RGB"
    working_image.convert(target_mode).save(buffer, format="PNG")
    buffer.seek(0)
    with Image.open(buffer) as png_image:
        return png_image.convert("RGBA").copy()


def make_thumbnail_image(image: Optional[Image.Image], size: Tuple[int, int] = (256, 256), *, fill_color=(255, 255, 255, 0)) -> Optional[Image.Image]:
    """Create a thumbnail copy of the provided image for gallery previews."""
    base_image = ensure_png_image(image)
    if base_image is None:
        return None

    preview = base_image.copy()
    try:
        preview.thumbnail(size, Image.Resampling.LANCZOS)
    except AttributeError:
        preview.thumbnail(size)
    canvas = Image.new("RGBA", size, fill_color)
    offset = ((size[0] - preview.width) // 2, (size[1] - preview.height) // 2)
    canvas.paste(preview, offset, preview if preview.mode == "RGBA" else None)
    return canvas


def add_to_history(image: Optional[Image.Image], workspace: str) -> Optional[str]:
    """Append processed image to history gallery with descriptive PNG name."""
    if image is None:
        return None

    png_image = ensure_png_image(image)
    if png_image is None:
        return None

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    prefix_map = {
        "bg_remove": "background_removed",
        "bg_gen": "background_generated",
        "face_swap": "face_swap",
        "video": "video_generated",
    }
    prefix = prefix_map.get(workspace, "result")
    file_name = f"{prefix}_{timestamp}.png"

    output_history.append((png_image, file_name))
    if len(output_history) > 20:
        del output_history[:-20]

    return file_name


def get_history_items(limit: Optional[int] = None):
    """Return history items (image, filename) with newest first."""
    if limit is None:
        data = output_history[:]
    elif limit <= 0:
        return []
    else:
        data = output_history[-limit:]
    return list(reversed(data))



class ImageProcessor:
    """Wrapper around PixelFree backend APIs."""

    def __init__(self) -> None:
        self.session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            pool_connections=10,
            pool_maxsize=20,
            max_retries=3,
        )
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)
        self.session.headers.update({"accept": "application/json"})

    def image_to_bytes(self, image: Image.Image, format: str = "PNG", preserve_alpha: bool = False) -> bytes:
        buffer = io.BytesIO()
        img = image
        if not preserve_alpha and img.mode in ("RGBA", "P"):
            img = img.convert("RGB")
        elif preserve_alpha and img.mode != "RGBA":
            img = img.convert("RGBA")
        img.save(buffer, format=format)
        buffer.seek(0)
        return buffer.getvalue()

    def base64_to_image(self, base64_str: str) -> Optional[Image.Image]:
        try:
            if base64_str.startswith("http"):
                logger.error("Refusing to download image from URL: %s", base64_str)
                return None
            if "," in base64_str:
                base64_str = base64_str.split(",", 1)[1]
            img_data = base64.b64decode(base64_str)
            return Image.open(io.BytesIO(img_data))
        except Exception as exc:
            logger.error("Failed to decode image: %s", exc)
            return None

    def call_api(self, endpoint: Optional[str], files: dict, data: dict) -> Optional[dict]:
        if not endpoint or not API_BASE_URL:
            logger.error("API endpoint or base URL missing")
            return None
        url = f"{API_BASE_URL}{endpoint}"
        try:
            response = self.session.post(url, files=files, data=data, timeout=300)
            logger.info("API response %s: %s", endpoint, response.status_code)
            response.raise_for_status()
            return response.json()
        except Exception as exc:
            logger.error("API call failed (%s): %s", endpoint, exc)
            return None

    def remove_background(self, image: Image.Image) -> Tuple[Optional[Image.Image], str]:
        start_time = time.time()
        files = {
            "file": ("input.png", self.image_to_bytes(image, preserve_alpha=True), "image/png")
        }
        result = self.call_api(API_ENDPOINTS.get("background_removal"), files, {})
        total = time.time() - start_time
        if result and result.get("success") and result.get("data"):
            b64 = result["data"].get("result_base64")
            image_obj = self.base64_to_image(b64) if b64 else None
            if image_obj:
                return image_obj, f"✅ 배경 제거 완료! (처리 시간: {total:.1f}초)"
        return None, f"❌ 배경 제거 실패 (시도 시간: {total:.1f}초)"

    def generate_background(
        self,
        image: Image.Image,
        bg_image: Image.Image,
        hq_mode: bool = False,
    ) -> Tuple[Optional[Image.Image], str]:
        start_time = time.time()
        files = {
            "image": ("input.png", self.image_to_bytes(image, preserve_alpha=True), "image/png"),
            "bg_image": ("background.png", self.image_to_bytes(bg_image), "image/png"),
        }
        data = {
            "jobId": f"job_{int(time.time() * 1000)}",
            "hq_mode": "true" if hq_mode else "false",
        }
        result = self.call_api(API_ENDPOINTS.get("background_generation"), files, data)
        total = time.time() - start_time
        mode_text = "고품질" if hq_mode else "일반"
        if result and result.get("success") and result.get("data"):
            payload = result["data"]
            if payload.get("imageData"):
                image_obj = self.base64_to_image(payload["imageData"])
                if image_obj:
                    return image_obj, f"✅ 배경 생성 완료! (처리 시간: {total:.1f}초, 모드: {mode_text})"
            if payload.get("rawDataUrl"):
                logger.error("Ignoring rawDataUrl field for security reasons.")
        return None, f"❌ 배경 생성 실패 (시도 시간: {total:.1f}초, 모드: {mode_text})"

    def swap_face(self, source: Image.Image, target: Image.Image) -> Tuple[Optional[Image.Image], str]:
        start_time = time.time()
        files = {
            "image": ("image.png", self.image_to_bytes(source), "image/png"),
            "targetFaceImage": ("targetFace.png", self.image_to_bytes(target), "image/png"),
        }
        data = {
            "jobId": f"face_swap_{int(time.time() * 1000)}"
        }
        result = self.call_api(API_ENDPOINTS.get("face_swap"), files, data)
        total = time.time() - start_time
        if result and result.get("success") and result.get("data"):
            payload = result["data"]
            if payload.get("imageData"):
                image_obj = self.base64_to_image(payload["imageData"])
                if image_obj:
                    api_time = payload.get("processingTime", 0)
                    return image_obj, f"✅ 얼굴 교체 완료! (전체: {total:.1f}초, API: {api_time:.1f}초)"
            if payload.get("rawDataUrl"):
                logger.error("Ignoring rawDataUrl field for security reasons.")
        return None, f"❌ 얼굴 교체 실패 (시도 시간: {total:.1f}초)"

    def generate_video(self, image: Optional[Image.Image]) -> Tuple[Optional[str], str]:
        if image is None:
            return None, "입력 이미지가 필요합니다."

        start_time = time.time()
        files = {
            "image": ("image.png", self.image_to_bytes(image, preserve_alpha=True), "image/png"),
        }
        data = {
            "jobId": f"video_{int(time.time() * 1000)}",
            "prompt": VIDEO_PROMPT,
            "duration": "5",
        }
        result = self.call_api(API_ENDPOINTS.get("video_generation"), files, data)
        total = time.time() - start_time
        if result and result.get("success") and result.get("data"):
            payload = result["data"]
            video_base64 = payload.get("videoData")
            if video_base64:
                if "," in video_base64:
                    video_base64 = video_base64.split(",", 1)[1]
                try:
                    video_bytes = base64.b64decode(video_base64)
                    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
                        temp_video.write(video_bytes)
                        temp_path = temp_video.name
                except Exception as exc:
                    logger.error("Failed to decode video: %s", exc)
                else:
                    api_time = payload.get("processingTime", 0)
                    return temp_path, f"✅ 영상 생성 완료! (전체: {total:.1f}초, API: {api_time:.1f}초)"
            if payload.get("rawDataUrl"):
                logger.error("Ignoring rawDataUrl field for security reasons.")
        return None, f"❌ 영상 생성 실패 (시도 시간: {total:.1f}초)"

processor = ImageProcessor()

# 인증 함수
def auth_fn(username: str, password: str) -> bool:
    """
    사용자 인증 함수
    - username이 허용된 사용자 목록에 있고
    - username과 password가 같으면 인증 성공
    """
    return username == password and username in ALLOWED_USERS

# 배경 이미지 로드
def load_background_images():
    """배경 이미지를 assets 폴더에서 로드 (Lazy loading)"""
    backgrounds = {
        "indoor": [],
        "outdoor": [],
        "indoor_paths": [],
        "outdoor_paths": [],
        "indoor_pil": [],
        "outdoor_pil": []
    }

    cwd = os.getcwd()
    
    # Indoor 이미지 경로 저장
    indoor_path = os.path.join(cwd, "assets", "indoor")
    if os.path.exists(indoor_path):
        files = sorted([f for f in os.listdir(indoor_path) if f.endswith('.png')])
        for img_file in files:
            img_path = os.path.join(indoor_path, img_file)
            try:
                backgrounds["indoor_paths"].append(img_path)
                # 갤러리용 썸네일만 생성
                with Image.open(img_path) as img:
                    img = img.convert('RGBA' if 'A' in img.getbands() else 'RGB')
                    thumb = img.copy()
                    thumb.thumbnail((240, 240), Image.Resampling.LANCZOS)
                backgrounds["indoor"].append(thumb)
                backgrounds["indoor_pil"].append(img_path)  # 경로 저장
            except Exception:
                pass
    
    # Outdoor 이미지 로드
    outdoor_path = os.path.join(cwd, "assets", "outdoor")
    if os.path.exists(outdoor_path):
        files = sorted([f for f in os.listdir(outdoor_path) if f.endswith('.png')])
        for img_file in files:
            img_path = os.path.join(outdoor_path, img_file)
            try:
                backgrounds["outdoor_paths"].append(img_path)
                with Image.open(img_path) as img:
                    img = img.convert('RGBA' if 'A' in img.getbands() else 'RGB')
                    thumb = img.copy()
                    thumb.thumbnail((240, 240), Image.Resampling.LANCZOS)
                backgrounds["outdoor"].append(thumb)
                backgrounds["outdoor_pil"].append(img_path)
            except Exception:
                pass
    
    # 이미지가 없는 경우 더미 이미지 생성
    if len(backgrounds['indoor']) == 0:
        dummy = Image.new('RGB', (240, 240), color=(200, 200, 200))
        backgrounds["indoor"].append(dummy)
        backgrounds["indoor_pil"].append(None)

    if len(backgrounds['outdoor']) == 0:
        dummy = Image.new('RGB', (240, 240), color=(150, 200, 150))
        backgrounds["outdoor"].append(dummy)
        backgrounds["outdoor_pil"].append(None)

    return backgrounds

bg_images = load_background_images()

# 얼굴 이미지 로드
def load_face_images():
    """얼굴 이미지들을 로드 - PIL 이미지로 직접 로드하여 갤러리에 전달"""
    faces = {
        "man_pil": [],
        "woman_pil": [],
        "all_pil": [],
        "man_paths": [],
        "woman_paths": [],
        "all_paths": []
    }
    
    cwd = os.getcwd()
    
    # Man 이미지 로드
    man_path = os.path.join(cwd, "assets", "faces", "man")
    if os.path.exists(man_path):
        files = sorted([f for f in os.listdir(man_path) if f.endswith(('.png', '.jpg', '.jpeg'))])
        for img_file in files:
            img_path = os.path.join(man_path, img_file)
            try:
                with Image.open(img_path) as img:
                    img = img.convert('RGBA' if 'A' in img.getbands() else 'RGB')
                    thumb = img.copy()
                    thumb.thumbnail((240, 240), Image.Resampling.LANCZOS)
                faces["man_pil"].append(thumb)
                faces["man_paths"].append(img_path)
            except Exception:
                pass
    
    # Woman 이미지 로드
    woman_path = os.path.join(cwd, "assets", "faces", "woman")
    if os.path.exists(woman_path):
        files = sorted([f for f in os.listdir(woman_path) if f.endswith(('.png', '.jpg', '.jpeg'))])
        for img_file in files:
            img_path = os.path.join(woman_path, img_file)
            try:
                with Image.open(img_path) as img:
                    img = img.convert('RGBA' if 'A' in img.getbands() else 'RGB')
                    thumb = img.copy()
                    thumb.thumbnail((240, 240), Image.Resampling.LANCZOS)
                faces["woman_pil"].append(thumb)
                faces["woman_paths"].append(img_path)
            except Exception:
                pass
    
    # 이미지가 없는 경우 더미 이미지 생성
    if len(faces['man_pil']) == 0:
        for _ in range(3):
            dummy = Image.new('RGB', (240, 240), color=(200, 200, 250))
            faces["man_pil"].append(dummy)
            faces["man_paths"].append(None)

    if len(faces['woman_pil']) == 0:
        for _ in range(3):
            dummy = Image.new('RGB', (240, 240), color=(250, 200, 200))
            faces["woman_pil"].append(dummy)
            faces["woman_paths"].append(None)

    # all 리스트 재구성
    faces["all_pil"] = faces["man_pil"] + faces["woman_pil"]
    faces["all_paths"] = faces["man_paths"] + faces["woman_paths"]
    
    return faces

face_images = load_face_images()
logger.info(f"Face images loaded - Total: {len(face_images.get('all_pil', []))}, Man: {len(face_images.get('man_pil', []))}, Woman: {len(face_images.get('woman_pil', []))}")

# 이미지를 실제로 로드하는 헬퍼 함수
def load_full_image(path_or_image):
    """경로에서 전체 크기 이미지를 로드 (캐싱 포함)"""
    if isinstance(path_or_image, str):
        cached = image_cache.get(path_or_image)
        if cached is not None:
            image_cache.move_to_end(path_or_image)
            return cached

        try:
            with Image.open(path_or_image) as img:
                mode = 'RGBA' if 'A' in img.getbands() else 'RGB'
                loaded = img.convert(mode).copy()
        except Exception as exc:
            logger.error(f"Failed to load image {path_or_image}: {exc}")
            return None

        cache_image(path_or_image, loaded)
        return loaded
    else:
        return path_or_image


def create_app():
    """PixelFree Gradio UI (v3.0.0)."""

    # CSS for history gallery image display
    custom_css = """
    #history-gallery img {
        object-fit: contain !important;
        height: 100% !important;
        width: 100% !important;
    }

    #history-gallery .thumbnail-item {
        height: 150px !important;
    }

    #history-gallery .grid-wrap {
        gap: 10px !important;
    }
    """

    with gr.Blocks(title="PixelFree v3.0.0", theme=gr.themes.Soft(), css=custom_css) as demo:

        # 세션 상태 관리
        session_user = gr.State(None)

        # 로그인 페이지
        with gr.Column(visible=True, elem_id="login-page") as login_page:
            gr.Markdown("# 🔐 PixelFree 로그인")
            gr.Markdown("허가된 사용자만 접근 가능합니다.")
            username_input = gr.Textbox(
                label="사용자명",
                placeholder="username을 입력하세요",
                max_lines=1
            )
            password_input = gr.Textbox(
                label="비밀번호",
                placeholder="password를 입력하세요",
                type="password",
                max_lines=1
            )
            login_btn = gr.Button("🔓 로그인", variant="primary", size="lg")
            login_msg = gr.Markdown("")

        # 메인 페이지 (초기에는 숨김)
        with gr.Column(visible=False, elem_id="main-page") as main_page:

            # 메인 헤더
            with gr.Row():
                gr.Markdown("# 🎨 PixelFree Studio")
                logout_btn = gr.Button("🚪 로그아웃", variant="secondary", size="sm")

            # 입력 섹션
            gr.Markdown("## 📸 입력 이미지")

            input_state = gr.State(None)
            bg_removed_state = gr.State(None)
            selected_bg_state = gr.State(None)
            selected_face_state = gr.State(None)
            last_output_state = gr.State(None)

            input_image = gr.Image(label="입력 이미지", type="pil", height=540)

            # 결과 섹션
            gr.Markdown("## ✨ 결과 이미지")
            output_image = gr.Image(label="처리된 이미지", type="pil", height=500, interactive=False)
            download_btn = gr.DownloadButton(
                "📥 사진 다운로드",
                value=None,
                interactive=False
            )

            # 비디오 생성 섹션
            gr.Markdown("## 🎬 AI 비디오 생성")
            with gr.Row():
                video_output = gr.Video(
                    label="생성된 영상",
                    height=360,
                    interactive=False,
                    autoplay=True,
                )
                with gr.Column():
                    gr.Markdown("최근 이미지 결과를 기반으로 짧은 클립을 생성합니다.")
                    generate_video_btn = gr.Button(
                        "🎬 영상 생성하기",
                        variant="primary",
                        size="lg",
                        interactive=False,
                    )

            # 배경 제거 섹션
            gr.Markdown("## 🎯 AI 배경 제거")
            remove_bg_btn = gr.Button("🚀 배경 제거 실행", variant="primary", size="lg", interactive=False)

            # AI 생성 기능 섹션
            gr.Markdown("## 🎨 AI 생성 기능")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 🏞️ 배경 생성")
                    with gr.Tabs():
                        with gr.Tab("🏠 실내"):
                            indoor_gallery = gr.Gallery(
                                value=bg_images["indoor"],
                                columns=4,
                                show_label=False, interactive=False,
                                allow_preview=False,
                                object_fit="contain"
                            )
                        with gr.Tab("🌳 아웃도어"):
                            outdoor_gallery = gr.Gallery(
                                value=bg_images["outdoor"],
                                columns=4,
                                rows=4, show_label=False, interactive=False,
                                allow_preview=False,
                                object_fit="contain"
                            )
                    hq_checkbox = gr.Checkbox(label="✨ 고품질 모드", value=False)
                    generate_bg_btn = gr.Button("🎨 배경 생성하기", variant="secondary", interactive=False, size="lg")
                with gr.Column():
                    gr.Markdown("### 👤 얼굴 교체")
                    with gr.Tabs():
                        with gr.Tab("👨 남성"):
                            man_gallery = gr.Gallery(
                                value=face_images["man_pil"],
                                columns=4,
                                rows=2, show_label=False, interactive=False,
                                allow_preview=False,
                                object_fit="contain"
                            )
                        with gr.Tab("👩 여성"):
                            woman_gallery = gr.Gallery(
                                value=face_images["woman_pil"],
                                columns=4,
                                rows=2, show_label=False, interactive=False,
                                allow_preview=False,
                                object_fit="contain"
                            )
                    swap_face_btn = gr.Button("🔄 얼굴 교체 적용", variant="secondary", interactive=False, size="lg")
            # Collapsible history section with image grid
            with gr.Accordion("📜 작업 히스토리", open=False):
                # History gallery with preview enabled
                history_gallery = gr.Gallery(
                    value=[img for img, _ in get_history_items(20)],
                    label=None,
                    show_label=False,
                    columns=5,
                    rows=4,
                    height=500,
                    object_fit="contain",
                    allow_preview=True,
                    elem_id="history-gallery",
                    interactive=False
                )

            def make_download_update(image: Optional[Image.Image], prefix: str = "pixelfree_output"):
                if image is None:
                    return gr.update(value=None, interactive=False)
                filename = f"{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
                temp_path = f"/tmp/{filename}"
                image.save(temp_path, "PNG")
                return gr.update(value=temp_path, interactive=True)

            def history_update():
                """Update history gallery with original images."""
                items = get_history_items(20)
                images = [img for img, _ in items]
                return gr.update(value=images)

            def disable_all_buttons():
                disabled = gr.update(interactive=False)
                return disabled, disabled, disabled, disabled

            def compute_button_states(
                input_img: Optional[Image.Image],
                last_output: Optional[Image.Image],
                bg_removed: Optional[Image.Image],
                selected_bg: Optional[Image.Image],
                selected_face: Optional[Image.Image],
            ):
                remove_ready = gr.update(interactive=input_img is not None)
                video_ready = gr.update(interactive=input_img is not None)
                bg_ready = gr.update(interactive=(bg_removed is not None and selected_bg is not None))
                swap_ready = gr.update(
                    interactive=(selected_face is not None and ((last_output is not None) or (input_img is not None)))
                )
                return remove_ready, video_ready, bg_ready, swap_ready


            def on_upload(new_image: Optional[Image.Image]):
                if new_image is None:
                    gr.Warning("입력 이미지를 업로드해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = disable_all_buttons()
                    return (
                        None,
                        None,
                        None,
                        None,
                        None,
                        gr.update(value=None),
                        make_download_update(None),
                        gr.update(value=None),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                        history_update(),
                    )

                gr.Info("입력 이미지를 불러왔습니다.")
                remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                    new_image,
                    None,
                    None,
                    None,
                    None,
                )
                return (
                    new_image,
                    None,
                    None,
                    None,
                    None,
                    gr.update(value=None),
                    make_download_update(None),
                    gr.update(value=None),
                    remove_ready,
                    video_ready,
                    bg_ready,
                    swap_ready,
                    history_update(),
                )

            def handle_remove_background(
                input_img: Optional[Image.Image],
                current_bg_removed: Optional[Image.Image],
                selected_bg: Optional[Image.Image],
                current_output: Optional[Image.Image],
                selected_face: Optional[Image.Image],
            ):
                if input_img is None:
                    gr.Warning("입력 이미지를 먼저 업로드해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        current_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        current_bg_removed,
                        current_output,
                        gr.update(),
                        make_download_update(current_output),
                        bg_ready,
                        swap_ready,
                        remove_ready,
                        video_ready,
                        history_update(),
                    )
                    return

                remove_disabled, video_disabled, bg_disabled, swap_disabled = disable_all_buttons()
                yield (
                    current_bg_removed,
                    current_output,
                    gr.update(),
                    make_download_update(current_output),
                    bg_disabled,
                    swap_disabled,
                    remove_disabled,
                    video_disabled,
                    history_update(),
                )

                removed, message = processor.remove_background(input_img)
                if removed is None:
                    gr.Warning(message)
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        current_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        current_bg_removed,
                        current_output,
                        gr.update(value=current_output),
                        make_download_update(current_output),
                        bg_ready,
                        swap_ready,
                        remove_ready,
                        video_ready,
                        history_update(),
                    )
                    return

                gr.Info(message)
                add_to_history(removed, "bg_remove")
                remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                    input_img,
                    removed,
                    removed,
                    selected_bg,
                    selected_face,
                )
                yield (
                    removed,
                    removed,
                    gr.update(value=removed),
                    make_download_update(removed, "background_removed"),
                    bg_ready,
                    swap_ready,
                    remove_ready,
                    video_ready,
                    history_update(),
                )

            def on_generate_video(
                last_output: Optional[Image.Image],
                input_img: Optional[Image.Image],
                current_bg_removed: Optional[Image.Image],
                selected_bg: Optional[Image.Image],
                selected_face: Optional[Image.Image],
            ):
                base_image = last_output or input_img
                if base_image is None:
                    gr.Warning("입력 이미지를 먼저 업로드해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        gr.update(value=None),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                remove_disabled, video_disabled, bg_disabled, swap_disabled = disable_all_buttons()
                yield (
                    gr.update(),
                    history_update(),
                    remove_disabled,
                    video_disabled,
                    bg_disabled,
                    swap_disabled,
                )

                video_path, message = processor.generate_video(base_image)
                if video_path is None:
                    gr.Warning(message)
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        gr.update(value=None),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                gr.Info(message)
                add_to_history(base_image, "video")
                remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                    input_img,
                    last_output,
                    current_bg_removed,
                    selected_bg,
                    selected_face,
                )
                yield (
                    gr.update(value=video_path),
                    history_update(),
                    remove_ready,
                    video_ready,
                    bg_ready,
                    swap_ready,
                )

            def load_background_image(category: str, index: int) -> Optional[Image.Image]:
                key = f"{category}_pil"
                paths = bg_images.get(key, [])
                if 0 <= index < len(paths):
                    path = paths[index]
                    return load_full_image(path)
                return None

            def on_select_background(
                evt: gr.SelectData,
                category: str,
                bg_removed: Optional[Image.Image],
            ):
                index = getattr(evt, "index", None)
                if index is None:
                    gr.Warning("선택 정보를 확인할 수 없습니다.")
                    return None, gr.update(interactive=False)
                image = load_background_image(category, index)
                if image is None:
                    gr.Warning("선택한 배경 이미지를 불러오지 못했습니다.")
                    return None, gr.update(interactive=False)
                gr.Info("배경 이미지를 선택했습니다.")
                ready = gr.update(interactive=bg_removed is not None)
                return image, ready

            def on_generate_background(
                last_output: Optional[Image.Image],
                selected_bg: Optional[Image.Image],
                hq_mode: bool,
                input_img: Optional[Image.Image],
                current_bg_removed: Optional[Image.Image],
                selected_face: Optional[Image.Image],
            ):
                if last_output is None:
                    gr.Warning("먼저 배경 제거를 실행해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        last_output,
                        gr.update(value=last_output),
                        make_download_update(last_output),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                if selected_bg is None:
                    gr.Warning("사용할 배경 이미지를 선택해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        last_output,
                        gr.update(value=last_output),
                        make_download_update(last_output),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                remove_disabled, video_disabled, bg_disabled, swap_disabled = disable_all_buttons()
                yield (
                    last_output,
                    gr.update(value=last_output),
                    make_download_update(last_output),
                    history_update(),
                    remove_disabled,
                    video_disabled,
                    bg_disabled,
                    swap_disabled,
                )

                generated, message = processor.generate_background(last_output, selected_bg, hq_mode)
                if generated is None:
                    gr.Warning(message)
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        last_output,
                        gr.update(value=last_output),
                        make_download_update(last_output),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                gr.Info(message)
                add_to_history(generated, "bg_gen")
                remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                    input_img,
                    generated,
                    current_bg_removed,
                    selected_bg,
                    selected_face,
                )
                yield (
                    generated,
                    gr.update(value=generated),
                    make_download_update(generated, "background_generated"),
                    history_update(),
                    remove_ready,
                    video_ready,
                    bg_ready,
                    swap_ready,
                )

            def load_face_image(category: str, index: int) -> Optional[Image.Image]:
                key = f"{category}_paths"
                paths = face_images.get(key, [])
                if 0 <= index < len(paths):
                    path = paths[index]
                    return load_full_image(path)
                return None

            def on_select_face(
                evt: gr.SelectData,
                category: str,
                last_output: Optional[Image.Image],
                input_img: Optional[Image.Image],
            ):
                base_image = last_output or input_img
                index = getattr(evt, "index", None)
                if index is None:
                    gr.Warning("선택 정보를 확인할 수 없습니다.")
                    return None, gr.update(interactive=False)
                image = load_face_image(category, index)
                if image is None:
                    gr.Warning("선택한 얼굴 이미지를 불러오지 못했습니다.")
                    return None, gr.update(interactive=False)
                gr.Info("얼굴 이미지를 선택했습니다.")
                ready = gr.update(interactive=base_image is not None)
                return image, ready

            def on_swap_face(
                last_output: Optional[Image.Image],
                input_img: Optional[Image.Image],
                selected_face: Optional[Image.Image],
                current_bg_removed: Optional[Image.Image],
                selected_bg: Optional[Image.Image],
            ):
                base_image = last_output or input_img
                if base_image is None:
                    gr.Warning("입력 이미지를 먼저 업로드해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        last_output,
                        gr.update(),
                        make_download_update(last_output),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                if selected_face is None:
                    gr.Warning("사용할 얼굴 이미지를 선택해주세요.")
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        last_output,
                        gr.update(value=base_image),
                        make_download_update(last_output),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                remove_disabled, video_disabled, bg_disabled, swap_disabled = disable_all_buttons()
                yield (
                    last_output,
                    gr.update(value=base_image),
                    make_download_update(last_output),
                    history_update(),
                    remove_disabled,
                    video_disabled,
                    bg_disabled,
                    swap_disabled,
                )

                swapped, message = processor.swap_face(base_image, selected_face)
                if swapped is None:
                    gr.Warning(message)
                    remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                        input_img,
                        last_output,
                        current_bg_removed,
                        selected_bg,
                        selected_face,
                    )
                    yield (
                        last_output,
                        gr.update(value=base_image),
                        make_download_update(last_output),
                        history_update(),
                        remove_ready,
                        video_ready,
                        bg_ready,
                        swap_ready,
                    )
                    return

                gr.Info(message)
                add_to_history(swapped, "face_swap")
                remove_ready, video_ready, bg_ready, swap_ready = compute_button_states(
                    input_img,
                    swapped,
                    current_bg_removed,
                    selected_bg,
                    selected_face,
                )
                yield (
                    swapped,
                    gr.update(value=swapped),
                    make_download_update(swapped, "face_swap"),
                    history_update(),
                    remove_ready,
                    video_ready,
                    bg_ready,
                    swap_ready,
                )

            input_image.change(
                on_upload,
                inputs=[input_image],
                outputs=[
                    input_state,
                    bg_removed_state,
                    selected_bg_state,
                    selected_face_state,
                    last_output_state,
                    output_image,
                    download_btn,
                    video_output,
                    remove_bg_btn,
                    generate_video_btn,
                    generate_bg_btn,
                    swap_face_btn,
                    history_gallery,
                ],
            )

            remove_bg_btn.click(
                handle_remove_background,
                inputs=[
                    input_state,
                    bg_removed_state,
                    selected_bg_state,
                    last_output_state,
                    selected_face_state,
                ],
                outputs=[
                    bg_removed_state,
                    last_output_state,
                    output_image,
                    download_btn,
                    generate_bg_btn,
                    swap_face_btn,
                    remove_bg_btn,
                    generate_video_btn,
                    history_gallery,
                ],
            )

            generate_video_btn.click(
                on_generate_video,
                inputs=[
                    last_output_state,
                    input_state,
                    bg_removed_state,
                    selected_bg_state,
                    selected_face_state,
                ],
                outputs=[
                    video_output,
                    history_gallery,
                    remove_bg_btn,
                    generate_video_btn,
                    generate_bg_btn,
                    swap_face_btn,
                ],
            )

            indoor_gallery.select(
                on_select_background,
                inputs=[gr.State("indoor"), bg_removed_state],
                outputs=[selected_bg_state, generate_bg_btn],
            )

            outdoor_gallery.select(
                on_select_background,
                inputs=[gr.State("outdoor"), bg_removed_state],
                outputs=[selected_bg_state, generate_bg_btn],
            )

            generate_bg_btn.click(
                on_generate_background,
                inputs=[
                    last_output_state,
                    selected_bg_state,
                    hq_checkbox,
                    input_state,
                    bg_removed_state,
                    selected_face_state,
                ],
                outputs=[
                    last_output_state,
                    output_image,
                    download_btn,
                    history_gallery,
                    remove_bg_btn,
                    generate_video_btn,
                    generate_bg_btn,
                    swap_face_btn,
                ],
            )

            man_gallery.select(
                on_select_face,
                inputs=[gr.State("man"), last_output_state, input_state],
                outputs=[selected_face_state, swap_face_btn],
            )

            woman_gallery.select(
                on_select_face,
                inputs=[gr.State("woman"), last_output_state, input_state],
                outputs=[selected_face_state, swap_face_btn],
            )

            swap_face_btn.click(
                on_swap_face,
                inputs=[
                    last_output_state,
                    input_state,
                    selected_face_state,
                    bg_removed_state,
                    selected_bg_state,
                ],
                outputs=[
                    last_output_state,
                    output_image,
                    download_btn,
                    history_gallery,
                    remove_bg_btn,
                    generate_video_btn,
                    generate_bg_btn,
                    swap_face_btn,
                ],
            )

        # 로그인/로그아웃 핸들러
        def handle_login(username, password):
            """로그인 처리"""
            if username in ALLOWED_USERS and username == password:
                return (
                    gr.update(visible=False),  # 로그인 페이지 숨김
                    gr.update(visible=True),   # 메인 페이지 표시
                    username,                  # 세션에 사용자 저장
                    gr.update(value="✅ 로그인 성공!")
                )
            return (
                gr.update(visible=True),   # 로그인 페이지 유지
                gr.update(visible=False),  # 메인 페이지 숨김
                None,                      # 세션 초기화
                gr.update(value="❌ 로그인 실패: 사용자명 또는 비밀번호가 올바르지 않습니다.")
            )

        def handle_logout():
            """로그아웃 처리"""
            return (
                gr.update(visible=True),   # 로그인 페이지 표시
                gr.update(visible=False),  # 메인 페이지 숨김
                None,                      # 세션 초기화
                gr.update(value="")        # 메시지 초기화
            )

        # 로그인/로그아웃 이벤트 연결
        login_btn.click(
            handle_login,
            inputs=[username_input, password_input],
            outputs=[login_page, main_page, session_user, login_msg]
        )

        logout_btn.click(
            handle_logout,
            outputs=[login_page, main_page, session_user, login_msg]
        )


    return demo

if __name__ == "__main__":
    app = create_app()
    app.launch(
        share=True,
        show_error=True,
        ssr_mode=False
    )