app.py

#!/usr/bin/env python

import os
import re
import tempfile
import gc  # garbage collector 추가
from collections.abc import Iterator
from threading import Thread
import json
import requests
import cv2
import base64
import logging
import time
from urllib.parse import quote  # URL 인코딩 (필요 시 사용)

import gradio as gr
import spaces
import torch
from loguru import logger
from PIL import Image
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer

# CSV/TXT/PDF 분석
import pandas as pd
import PyPDF2

# =============================================================================
# (신규) 이미지 API 관련 함수들
# =============================================================================
from gradio_client import Client

API_URL = "http://211.233.58.201:7896"

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def test_api_connection() -> str:
    """API 서버 연결 테스트"""
    try:
        client = Client(API_URL)
        return "API 연결 성공: 정상 작동 중"
    except Exception as e:
        logging.error(f"API connection test failed: {e}")
        return f"API 연결 실패: {e}"

def generate_image(prompt: str, width: float, height: float, guidance: float, inference_steps: float, seed: float):
    """
    이미지 생성 함수.
    여기서는 서버가 최종 이미지를 Base64(또는 data:image/...) 형태로 직접 반환한다고 가정합니다.
    /tmp/... 경로나 추가 다운로드를 시도하지 않습니다.
    """
    if not prompt:
        return None, "Error: Prompt is required"
    try:
        logging.info(f"Calling image generation API with prompt: {prompt}")
        
        client = Client(API_URL)
        result = client.predict(
            prompt=prompt,
            width=int(width),
            height=int(height),
            guidance=float(guidance),
            inference_steps=int(inference_steps),
            seed=int(seed),
            do_img2img=False,
            init_image=None,
            image2image_strength=0.8,
            resize_img=True,
            api_name="/generate_image"
        )
        
        logging.info(
            f"Image generation result: {type(result)}, "
            f"length: {len(result) if isinstance(result, (list, tuple)) else 'unknown'}"
        )
        
        # 결과가 튜플/리스트: [이미지_base64 or data_url, seed_info] 로 가정
        if isinstance(result, (list, tuple)) and len(result) > 0:
            image_data = result[0]  # 첫 번째 요소가 이미지 데이터 (Base64 or data:image/... 등)
            seed_info = result[1] if len(result) > 1 else "Unknown seed"
            return image_data, seed_info
        else:
            # 다른 형태로 반환된 경우
            return result, "Unknown seed"
            
    except Exception as e:
        logging.error(f"Image generation failed: {str(e)}")
        return None, f"Error: {str(e)}"

# Base64 패딩 수정 함수 (필요하다면 사용)
def fix_base64_padding(data):
    """Base64 문자열의 패딩을 수정합니다."""
    if isinstance(data, bytes):
        data = data.decode('utf-8')
    
    if "base64," in data:
        data = data.split("base64,", 1)[1]
    
    missing_padding = len(data) % 4
    if missing_padding:
        data += '=' * (4 - missing_padding)
    
    return data

# =============================================================================
# 메모리 정리 함수
# =============================================================================
def clear_cuda_cache():
    """CUDA 캐시를 명시적으로 비웁니다."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()

# =============================================================================
# SerpHouse 관련 함수
# =============================================================================
SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "")

def extract_keywords(text: str, top_k: int = 5) -> str:
    """단순 키워드 추출: 한글, 영어, 숫자, 공백만 남김"""
    text = re.sub(r"[^a-zA-Z0-9가-힣\s]", "", text)
    tokens = text.split()
    return " ".join(tokens[:top_k])

def do_web_search(query: str) -> str:
    """
    SerpHouse LIVE API 호출하여 검색 결과 마크다운 반환
    (필요하다면 수정 or 삭제 가능)
    """
    try:
        url = "https://api.serphouse.com/serp/live"
        params = {
            "q": query,
            "domain": "google.com",
            "serp_type": "web",
            "device": "desktop",
            "lang": "en",
            "num": "20"
        }
        headers = {"Authorization": f"Bearer {SERPHOUSE_API_KEY}"}
        logger.info(f"SerpHouse API 호출 중... 검색어: {query}")
        response = requests.get(url, headers=headers, params=params, timeout=60)
        response.raise_for_status()
        data = response.json()
        results = data.get("results", {})
        organic = None
        if isinstance(results, dict) and "organic" in results:
            organic = results["organic"]
        elif isinstance(results, dict) and "results" in results:
            if isinstance(results["results"], dict) and "organic" in results["results"]:
                organic = results["results"]["organic"]
        elif "organic" in data:
            organic = data["organic"]
        if not organic:
            logger.warning("응답에서 organic 결과를 찾을 수 없습니다.")
            return "No web search results found or unexpected API response structure."
        max_results = min(20, len(organic))
        limited_organic = organic[:max_results]
        summary_lines = []
        for idx, item in enumerate(limited_organic, start=1):
            title = item.get("title", "No title")
            link = item.get("link", "#")
            snippet = item.get("snippet", "No description")
            displayed_link = item.get("displayed_link", link)
            summary_lines.append(
                f"### Result {idx}: {title}\n\n"
                f"{snippet}\n\n"
                f"**출처**: [{displayed_link}]({link})\n\n"
                f"---\n"
            )
        instructions = """
# 웹 검색 결과
아래는 검색 결과입니다. 질문에 답변할 때 이 정보를 활용하세요:
1. 여러 출처 내용을 종합하여 답변.
2. 출처 인용 시 "[출처 제목](링크)" 마크다운 형식 사용.
3. 답변 마지막에 '참고 자료:' 섹션에 사용한 주요 출처를 나열.
"""
        return instructions + "\n".join(summary_lines)
    except Exception as e:
        logger.error(f"Web search failed: {e}")
        return f"Web search failed: {str(e)}"

# =============================================================================
# 모델 및 프로세서 로딩
# =============================================================================
MAX_CONTENT_CHARS = 2000
MAX_INPUT_LENGTH = 2096

model_id = os.getenv("MODEL_ID", "VIDraft/Gemma-3-R1984-4B")
processor = AutoProcessor.from_pretrained(model_id, padding_side="left")
model = Gemma3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="eager"
)

MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))

# =============================================================================
# CSV, TXT, PDF 분석 함수
# =============================================================================
def analyze_csv_file(path: str) -> str:
    try:
        df = pd.read_csv(path)
        if df.shape[0] > 50 or df.shape[1] > 10:
            df = df.iloc[:50, :10]
        df_str = df.to_string()
        if len(df_str) > MAX_CONTENT_CHARS:
            df_str = df_str[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
        return f"**[CSV File: {os.path.basename(path)}]**\n\n{df_str}"
    except Exception as e:
        return f"Failed to read CSV ({os.path.basename(path)}): {str(e)}"

def analyze_txt_file(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            text = f.read()
        if len(text) > MAX_CONTENT_CHARS:
            text = text[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
        return f"**[TXT File: {os.path.basename(path)}]**\n\n{text}"
    except Exception as e:
        return f"Failed to read TXT ({os.path.basename(path)}): {str(e)}"

def pdf_to_markdown(pdf_path: str) -> str:
    text_chunks = []
    try:
        with open(pdf_path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            max_pages = min(5, len(reader.pages))
            for page_num in range(max_pages):
                page_text = reader.pages[page_num].extract_text() or ""
                page_text = page_text.strip()
                if page_text:
                    if len(page_text) > MAX_CONTENT_CHARS // max_pages:
                        page_text = page_text[:MAX_CONTENT_CHARS // max_pages] + "...(truncated)"
                    text_chunks.append(f"## Page {page_num+1}\n\n{page_text}\n")
            if len(reader.pages) > max_pages:
                text_chunks.append(f"\n...(Showing {max_pages} of {len(reader.pages)} pages)...")
    except Exception as e:
        return f"Failed to read PDF ({os.path.basename(pdf_path)}): {str(e)}"
    full_text = "\n".join(text_chunks)
    if len(full_text) > MAX_CONTENT_CHARS:
        full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
    return f"**[PDF File: {os.path.basename(pdf_path)}]**\n\n{full_text}"

# =============================================================================
# 이미지/비디오 파일 제한 검사
# =============================================================================
def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
    image_count = 0
    video_count = 0
    for path in paths:
        if path.endswith(".mp4"):
            video_count += 1
        elif re.search(r"\.(png|jpg|jpeg|gif|webp)$", path, re.IGNORECASE):
            image_count += 1
    return image_count, video_count

def count_files_in_history(history: list[dict]) -> tuple[int, int]:
    image_count = 0
    video_count = 0
    for item in history:
        if item["role"] != "user" or isinstance(item["content"], str):
            continue
        if isinstance(item["content"], list) and len(item["content"]) > 0:
            file_path = item["content"][0]
            if isinstance(file_path, str):
                if file_path.endswith(".mp4"):
                    video_count += 1
                elif re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE):
                    image_count += 1
    return image_count, video_count

def validate_media_constraints(message: dict, history: list[dict]) -> bool:
    """이미지/비디오 업로드 제한 검사."""
    media_files = [f for f in message["files"] 
                   if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE) or f.endswith(".mp4")]
    new_image_count, new_video_count = count_files_in_new_message(media_files)
    history_image_count, history_video_count = count_files_in_history(history)

    image_count = history_image_count + new_image_count
    video_count = history_video_count + new_video_count

    if video_count > 1:
        gr.Warning("Only one video is supported.")
        return False
    if video_count == 1:
        if image_count > 0:
            gr.Warning("Mixing images and videos is not allowed.")
            return False
        if "<image>" in message["text"]:
            gr.Warning("Using <image> tags with video files is not supported.")
            return False
    if video_count == 0 and image_count > MAX_NUM_IMAGES:
        gr.Warning(f"You can upload up to {MAX_NUM_IMAGES} images.")
        return False
    if "<image>" in message["text"]:
        image_files = [f for f in message["files"]
                       if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)]
        image_tag_count = message["text"].count("<image>")
        if image_tag_count != len(image_files):
            gr.Warning("The number of <image> tags in the text does not match the number of image files.")
            return False
    return True

# =============================================================================
# 비디오 처리 함수
# =============================================================================
def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
    vidcap = cv2.VideoCapture(video_path)
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_interval = max(int(fps), int(total_frames / 10))
    frames = []
    for i in range(0, total_frames, frame_interval):
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        success, image = vidcap.read()
        if success:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (0, 0), fx=0.5, fy=0.5)
            pil_image = Image.fromarray(image)
            timestamp = round(i / fps, 2)
            frames.append((pil_image, timestamp))
            if len(frames) >= 5:
                break
    vidcap.release()
    return frames

def process_video(video_path: str) -> tuple[list[dict], list[str]]:
    content = []
    temp_files = []
    frames = downsample_video(video_path)
    for pil_image, timestamp in frames:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
            pil_image.save(temp_file.name)
            temp_files.append(temp_file.name)
            content.append({"type": "text", "text": f"Frame {timestamp}:"})
            content.append({"type": "image", "url": temp_file.name})
    return content, temp_files

# =============================================================================
# interleaved <image> 처리 함수 (<image> 태그와 이미지 업로드 혼합 지원)
# =============================================================================
def process_interleaved_images(message: dict) -> list[dict]:
    parts = re.split(r"(<image>)", message["text"])
    content = []
    image_files = [f for f in message["files"]
                   if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)]
    image_index = 0
    for part in parts:
        if part == "<image>" and image_index < len(image_files):
            content.append({"type": "image", "url": image_files[image_index]})
            image_index += 1
        elif part.strip():
            content.append({"type": "text", "text": part.strip()})
        else:
            if isinstance(part, str) and part != "<image>":
                content.append({"type": "text", "text": part})
    return content

# =============================================================================
# 파일 처리 -> content 생성
# =============================================================================
def is_image_file(file_path: str) -> bool:
    return bool(re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE))

def is_video_file(file_path: str) -> bool:
    return file_path.endswith(".mp4")

def is_document_file(file_path: str) -> bool:
    return file_path.lower().endswith(".pdf") or file_path.lower().endswith(".csv") or file_path.lower().endswith(".txt")

def process_new_user_message(message: dict) -> tuple[list[dict], list[str]]:
    """사용자가 새로 입력한 메시지 + 업로드 파일들을 하나의 content(list)로 변환."""
    temp_files = []
    if not message["files"]:
        return [{"type": "text", "text": message["text"]}], temp_files

    video_files = [f for f in message["files"] if is_video_file(f)]
    image_files = [f for f in message["files"] if is_image_file(f)]
    csv_files = [f for f in message["files"] if f.lower().endswith(".csv")]
    txt_files = [f for f in message["files"] if f.lower().endswith(".txt")]
    pdf_files = [f for f in message["files"] if f.lower().endswith(".pdf")]

    content_list = [{"type": "text", "text": message["text"]}]

    # 문서들
    for csv_path in csv_files:
        content_list.append({"type": "text", "text": analyze_csv_file(csv_path)})
    for txt_path in txt_files:
        content_list.append({"type": "text", "text": analyze_txt_file(txt_path)})
    for pdf_path in pdf_files:
        content_list.append({"type": "text", "text": pdf_to_markdown(pdf_path)})

    # 비디오 처리
    if video_files:
        video_content, video_temp_files = process_video(video_files[0])
        content_list += video_content
        temp_files.extend(video_temp_files)
        return content_list, temp_files

    # 이미지 처리
    if "<image>" in message["text"] and image_files:
        interleaved_content = process_interleaved_images({"text": message["text"], "files": image_files})
        if content_list and content_list[0]["type"] == "text":
            content_list = content_list[1:]
        return interleaved_content + content_list, temp_files
    else:
        for img_path in image_files:
            content_list.append({"type": "image", "url": img_path})

    return content_list, temp_files

# =============================================================================
# history -> LLM 메시지 변환
# =============================================================================
def process_history(history: list[dict]) -> list[dict]:
    """
    기존 대화 기록을 LLM에 맞게 변환.
    - user -> {"role":"user","content":[{type,text},...]}
    - assistant -> {"role":"assistant","content":[{type:"text",text},...]}
    """
    messages = []
    current_user_content = []
    for item in history:
        if item["role"] == "assistant":
            # 사용자 content 누적분이 있으면 한번에 user로 추가
            if current_user_content:
                messages.append({"role": "user", "content": current_user_content})
                current_user_content = []
            # assistant 바로 추가
            messages.append({"role": "assistant", "content": [{"type": "text", "text": item["content"]}]})
        else:
            content = item["content"]
            if isinstance(content, str):
                current_user_content.append({"type": "text", "text": content})
            elif isinstance(content, list) and len(content) > 0:
                file_path = content[0]
                if is_image_file(file_path):
                    current_user_content.append({"type": "image", "url": file_path})
                else:
                    current_user_content.append({"type": "text", "text": f"[File: {os.path.basename(file_path)}]"})
    if current_user_content:
        messages.append({"role": "user", "content": current_user_content})
    return messages

# =============================================================================
# 모델 생성 함수 (OOM 캐치)
# =============================================================================
def _model_gen_with_oom_catch(**kwargs):
    try:
        model.generate(**kwargs)
    except torch.cuda.OutOfMemoryError:
        raise RuntimeError("[OutOfMemoryError] GPU 메모리가 부족합니다.")
    finally:
        clear_cuda_cache()

# =============================================================================
# 메인 추론 함수
# =============================================================================
@spaces.GPU(duration=120)
def run(
    message: dict,
    history: list[dict],
    system_prompt: str = "",
    max_new_tokens: int = 512,
    use_web_search: bool = False,
    web_search_query: str = "",
    age_group: str = "20대",
    mbti_personality: str = "INTP",
    sexual_openness: int = 2,
    image_gen: bool = False
) -> Iterator[str]:
    """
    LLM 추론 함수.
    - 이미지 생성 시, 서버가 Base64(또는 data:image/... 형태)를 직접 반환한다고 가정.
    - /tmp/... 파일에 대한 재다운로드를 시도하지 않음 (403 Forbidden 문제 회피).
    """
    if not validate_media_constraints(message, history):
        yield ""
        return

    temp_files = []
    try:
        # 1) 시스템 프롬프트 + 페르소나 정보
        persona = (
            f"{system_prompt.strip()}\n\n"
            f"Gender: Female\n"
            f"Age Group: {age_group}\n"
            f"MBTI Persona: {mbti_personality}\n"
            f"Sexual Openness (1~5): {sexual_openness}\n"
        )
        combined_system_msg = f"[System Prompt]\n{persona.strip()}\n\n"

        # 2) 웹 검색 (옵션)
        if use_web_search:
            user_text = message["text"]
            ws_query = extract_keywords(user_text)
            if ws_query.strip():
                logger.info(f"[Auto WebSearch Keyword] {ws_query!r}")
                ws_result = do_web_search(ws_query)
                combined_system_msg += f"[Search top-20 Full Items]\n{ws_result}\n\n"
                combined_system_msg += (
                    "[참고: 위 검색결과 link를 출처로 인용하여 답변]\n"
                    "[중요 지시사항]\n"
                    "1. 검색 결과에서 찾은 정보의 출처를 반드시 인용.\n"
                    "2. '[출처 제목](링크)' 형식으로 링크.\n"
                    "3. 답변 마지막에 '참고 자료:' 섹션.\n"
                )
            else:
                combined_system_msg += "[No valid keywords found, skipping WebSearch]\n\n"

        # 3) 기존 history + 새 user 메시지
        messages = []
        if combined_system_msg.strip():
            messages.append({"role": "system", "content": [{"type": "text", "text": combined_system_msg.strip()}]})
        messages.extend(process_history(history))

        user_content, user_temp_files = process_new_user_message(message)
        temp_files.extend(user_temp_files)

        for item in user_content:
            if item["type"] == "text" and len(item["text"]) > MAX_CONTENT_CHARS:
                item["text"] = item["text"][:MAX_CONTENT_CHARS] + "\n...(truncated)..."

        messages.append({"role": "user", "content": user_content})

        # 4) 토크나이징
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(device=model.device, dtype=torch.bfloat16)
        if inputs.input_ids.shape[1] > MAX_INPUT_LENGTH:
            inputs.input_ids = inputs.input_ids[:, -MAX_INPUT_LENGTH:]
            if 'attention_mask' in inputs:
                inputs.attention_mask = inputs.attention_mask[:, -MAX_INPUT_LENGTH:]

        streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)

        t = Thread(target=_model_gen_with_oom_catch, kwargs=gen_kwargs)
        t.start()

        # 스트리밍 출력
        output_so_far = ""
        for new_text in streamer:
            output_so_far += new_text
            yield output_so_far

        # 5) 이미지 생성 (Base64)
        if image_gen:
            last_user_text = message["text"].strip()
            if not last_user_text:
                yield output_so_far + "\n\n(이미지 생성 실패: Empty user prompt)"
            else:
                try:
                    width, height = 512, 512
                    guidance, steps, seed = 7.5, 30, 42
                    
                    logger.info(f"Generating image with prompt: {last_user_text}")
                    
                    # API 호출해서 (base64) 이미지 생성
                    image_result, seed_info = generate_image(
                        prompt=last_user_text,
                        width=width, 
                        height=height, 
                        guidance=guidance, 
                        inference_steps=steps, 
                        seed=seed
                    )
                    
                    logger.info(f"Received image data type: {type(image_result)}")

                    # Base64 or data:image/... 처리
                    if image_result:
                        if isinstance(image_result, str):
                            # 이미 data:image/로 시작하면 그대로 사용
                            if image_result.startswith("data:image/"):
                                final_md = f"\n\n**[생성된 이미지]**\n\n![생성된 이미지]({image_result})"
                                yield output_so_far + final_md
                            else:
                                # 순수 base64로 판단(단, 일반 URL이나 '/tmp/...'이면 처리 불가)
                                if len(image_result) > 100 and "/" not in image_result:
                                    # base64
                                    image_data = "data:image/webp;base64," + image_result
                                    final_md = f"\n\n**[생성된 이미지]**\n\n![생성된 이미지]({image_data})"
                                    yield output_so_far + final_md
                                else:
                                    # 그 외 (ex. http://..., /tmp/...) -> 403 문제 발생하므로 표시 안 함
                                    yield output_so_far + "\n\n(이미지 생성 결과가 base64 형식이 아닙니다)"
                        else:
                            yield output_so_far + "\n\n(이미지 생성 결과가 문자열이 아님)"
                    else:
                        yield output_so_far + f"\n\n(이미지 생성 실패: {seed_info})"

                except Exception as e:
                    logger.error(f"Image generation error: {e}")
                    yield output_so_far + f"\n\n(이미지 생성 중 오류 발생: {e})"

    except Exception as e:
        logger.error(f"Error in run: {str(e)}")
        yield f"죄송합니다. 오류가 발생했습니다: {str(e)}"
    finally:
        for tmp in temp_files:
            try:
                if os.path.exists(tmp):
                    os.unlink(tmp)
                    logger.info(f"Deleted temp file: {tmp}")
            except Exception as ee:
                logger.warning(f"Failed to delete temp file {tmp}: {ee}")
        try:
            del inputs, streamer
        except Exception:
            pass
        clear_cuda_cache()

# =============================================================================
# 예시들
# =============================================================================
examples = [
    [
        {
            "text": "Compare the contents of the two PDF files.",
            "files": [
                "assets/additional-examples/before.pdf",
                "assets/additional-examples/after.pdf",
            ],
        }
    ],
    [
        {
            "text": "Summarize and analyze the contents of the CSV file.",
            "files": ["assets/additional-examples/sample-csv.csv"],
        }
    ],
    # ... 나머지 예시 필요하다면 추가 ...
]

# =============================================================================
# Gradio UI (Blocks) 구성
# =============================================================================

css = """
.gradio-container {
    background: rgba(255, 255, 255, 0.7);
    padding: 30px 40px;
    margin: 20px auto;
    width: 100% !important;
    max-width: none !important;
}
"""
title_html = """
<h1 align="center" style="margin-bottom: 0.2em; font-size: 1.6em;"> 💘 HeartSync : Love Dating AI 💘 </h1>
<p align="center" style="font-size:1.1em; color:#555;">
    ✅ FLUX Image Generation ✅ Reasoning & Uncensored ✅ Multimodal & VLM ✅ Deep-Research & RAG <br>
</p>
"""

with gr.Blocks(css=css, title="HeartSync") as demo:
    gr.Markdown(title_html)
    
    # 별도 갤러리 예시 (필요 시 사용)
    generated_images = gr.Gallery(
        label="생성된 이미지", 
        show_label=True, 
        visible=False,
        elem_id="generated_images",
        columns=2,
        height="auto",
        object_fit="contain"
    )
    
    with gr.Row():
        web_search_checkbox = gr.Checkbox(label="Deep Research", value=False)
        image_gen_checkbox = gr.Checkbox(label="Image Gen", value=False)
    
    base_system_prompt_box = gr.Textbox(
        lines=3,
        value="You are a deep thinking AI...\n페르소나: 당신은 달콤하고...",
        label="기본 시스템 프롬프트",
        visible=False
    )
    with gr.Row():
        age_group_dropdown = gr.Dropdown(
            label="연령대 선택 (기본 20대)",
            choices=["10대", "20대", "30~40대", "50~60대", "70대 이상"],
            value="20대",
            interactive=True
        )
    mbti_choices = [
        "INTJ (용의주도한 전략가)",
        "INTP (논리적인 사색가)",
        "ENTJ (대담한 통솔자)",
        "ENTP (뜨거운 논쟁가)",
        "INFJ (선의의 옹호자)",
        "INFP (열정적인 중재자)",
        "ENFJ (정의로운 사회운동가)",
        "ENFP (재기발랄한 활동가)",
        "ISTJ (청렴결백한 논리주의자)",
        "ISFJ (용감한 수호자)",
        "ESTJ (엄격한 관리자)",
        "ESFJ (사교적인 외교관)",
        "ISTP (만능 재주꾼)",
        "ISFP (호기심 많은 예술가)",
        "ESTP (모험을 즐기는 사업가)",
        "ESFP (자유로운 영혼의 연예인)"
    ]
    mbti_dropdown = gr.Dropdown(
        label="AI 페르소나 MBTI (기본 INTP)",
        choices=mbti_choices,
        value="INTP (논리적인 사색가)",
        interactive=True
    )
    sexual_openness_slider = gr.Slider(
        minimum=1, maximum=5, step=1, value=2,
        label="섹슈얼 관심도/개방성 (1~5, 기본=2)",
        interactive=True
    )
    max_tokens_slider = gr.Slider(
        label="Max New Tokens",
        minimum=100, maximum=8000, step=50, value=1000,
        visible=False
    )
    web_search_text = gr.Textbox(
        lines=1,
        label="(Unused) Web Search Query",
        placeholder="No direct input needed",
        visible=False
    )

    def modified_run(
        message, history, system_prompt, max_new_tokens, 
        use_web_search, web_search_query,
        age_group, mbti_personality, sexual_openness, image_gen
    ):
        """
        run() 함수를 호출하여 텍스트 스트림을 받고, 
        필요 시 추가 처리 후 결과 반환 (갤러리 업데이트 등).
        """
        output_so_far = ""
        gallery_update = gr.Gallery(visible=False, value=[])
        yield output_so_far, gallery_update

        text_generator = run(
            message, history, 
            system_prompt, max_new_tokens,
            use_web_search, web_search_query,
            age_group, mbti_personality,
            sexual_openness, image_gen
        )
        
        for text_chunk in text_generator:
            output_so_far = text_chunk
            yield output_so_far, gallery_update

        # 만약 run() 내부에서 Base64 이미지를 이미 대화창에 삽입했다면,
        # 여기서 갤러리에 따로 표시할 필요는 없을 수도 있습니다.
        # run() 내부에서의 image_result를 가져오려면, run() 함수가 해당 정보를 반환하도록 추가 수정이 필요합니다.

    chat = gr.ChatInterface(
        fn=modified_run,
        type="messages",
        chatbot=gr.Chatbot(type="messages", scale=1, allow_tags=["image"]),
        textbox=gr.MultimodalTextbox(
            file_types=[".webp", ".png", ".jpg", ".jpeg", ".gif", ".mp4", ".csv", ".txt", ".pdf"],
            file_count="multiple",
            autofocus=True
        ),
        multimodal=True,
        additional_inputs=[
            base_system_prompt_box,
            max_tokens_slider,
            web_search_checkbox,
            web_search_text,
            age_group_dropdown,
            mbti_dropdown,
            sexual_openness_slider,
            image_gen_checkbox,
        ],
        additional_outputs=[generated_images],
        stop_btn=False,
        title='<a href="https://discord.gg/openfreeai" target="_blank">https://discord.gg/openfreeai</a>',
        examples=examples,
        run_examples_on_click=False,
        cache_examples=False,
        css_paths=None,
        delete_cache=(1800, 1800),
    )

    with gr.Row(elem_id="examples_row"):
        with gr.Column(scale=12, elem_id="examples_container"):
            gr.Markdown("### Example Inputs (click to load)")

if __name__ == "__main__":
    demo.launch(share=True)