Update app.py

app.py

@@ -9,6 +9,11 @@ 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
@@ -16,162 +21,189 @@ from loguru import logger
 from PIL import Image
 from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer
 
-# CSV/TXT 분석
+# CSV/TXT/PDF 분석
 import pandas as pd
-# PDF 텍스트 추출
 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 API key from environment variable
-##############################################################################
+# =============================================================================
+# SerpHouse 관련 함수
+# =============================================================================
 SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "")
 
-##############################################################################
-# 간단한 키워드 추출 함수 (한글 + 알파벳 + 숫자 + 공백 보존)
-##############################################################################
 def extract_keywords(text: str, top_k: int = 5) -> str:
-    """
-    1) 한글(가-힣), 영어(a-zA-Z), 숫자(0-9), 공백만 남김
-    2) 공백 기준 토큰 분리
-    3) 최대 top_k개만
-    """
+    """단순 키워드 추출: 한글, 영어, 숫자, 공백만 남김"""
     text = re.sub(r"[^a-zA-Z0-9가-힣\s]", "", text)
     tokens = text.split()
-    key_tokens = tokens[:top_k]
-    return " ".join(key_tokens)
+    return " ".join(tokens[:top_k])
 
-##############################################################################
-# SerpHouse Live endpoint 호출
-# - 상위 20개 결과 JSON을 LLM에 넘길 때 link, snippet 등 모두 포함
-##############################################################################
 def do_web_search(query: str) -> str:
     """
-    상위 20개 'organic' 결과 item 전체(제목, link, snippet 등)를
-    JSON 문자열 형태로 반환
+    SerpHouse LIVE API 호출하여 검색 결과 마크다운 반환
+    (필요하다면 수정 or 삭제 가능)
     """
     try:
         url = "https://api.serphouse.com/serp/live"
-        
-        # 기본 GET 방식으로 파라미터 간소화하고 결과 수를 20개로 제한
         params = {
             "q": query,
             "domain": "google.com",
-            "serp_type": "web",  # 기본 웹 검색
+            "serp_type": "web",
             "device": "desktop",
             "lang": "en",
-            "num": "20"  # 최대 20개 결과만 요청
-        }
-        
-        headers = {
-            "Authorization": f"Bearer {SERPHOUSE_API_KEY}"
+            "num": "20"
         }
-        
+        headers = {"Authorization": f"Bearer {SERPHOUSE_API_KEY}"}
         logger.info(f"SerpHouse API 호출 중... 검색어: {query}")
-        logger.info(f"요청 URL: {url} - 파라미터: {params}")
-        
-        # GET 요청 수행
         response = requests.get(url, headers=headers, params=params, timeout=60)
         response.raise_for_status()
-        
-        logger.info(f"SerpHouse API 응답 상태 코드: {response.status_code}")
         data = response.json()
-        
-        # 다양한 응답 구조 처리
         results = data.get("results", {})
         organic = None
-        
-        # 가능한 응답 구조 1
         if isinstance(results, dict) and "organic" in results:
             organic = results["organic"]
-        
-        # 가능한 응답 구조 2 (중첩된 results)
         elif isinstance(results, dict) and "results" in results:
             if isinstance(results["results"], dict) and "organic" in results["results"]:
                 organic = results["results"]["organic"]
-        
-        # 가능한 응답 구조 3 (최상위 organic)
         elif "organic" in data:
             organic = data["organic"]
-            
         if not organic:
             logger.warning("응답에서 organic 결과를 찾을 수 없습니다.")
-            logger.debug(f"응답 구조: {list(data.keys())}")
-            if isinstance(results, dict):
-                logger.debug(f"results 구조: {list(results.keys())}")
             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. 답변에 관련 정보의 출처를 명시적으로 인용하세요 (예: "X 출처에 따르면...")
-3. 응답에 실제 출처 링크를 포함하세요
-4. 여러 출처의 정보를 종합하여 답변하세요
+1. 여러 출처 내용을 종합하여 답변.
+2. 출처 인용 시 "[출처 제목](링크)" 마크다운 형식 사용.
+3. 답변 마지막에 '참고 자료:' 섹션에 사용한 주요 출처를 나열.
 """
-        
-        search_results = instructions + "\n".join(summary_lines)
-        logger.info(f"검색 결과 {len(limited_organic)}개 처리 완료")
-        return search_results
-    
+        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")
+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"  # 가능하다면 "flash_attention_2"로 변경
+    attn_implementation="eager"
 )
-MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
 
+MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
 
-##############################################################################
+# =============================================================================
 # CSV, TXT, PDF 분석 함수
-##############################################################################
+# =============================================================================
 def analyze_csv_file(path: str) -> str:
-    """
-    CSV 파일을 전체 문자열로 변환. 너무 길 경우 일부만 표시.
-    """
     try:
         df = pd.read_csv(path)
         if df.shape[0] > 50 or df.shape[1] > 10:
@@ -183,11 +215,7 @@ def analyze_csv_file(path: str) -> str:
     except Exception as e:
         return f"Failed to read CSV ({os.path.basename(path)}): {str(e)}"
 
-
 def analyze_txt_file(path: str) -> str:
-    """
-    TXT 파일 전문 읽기. 너무 길면 일부만 표시.
-    """
     try:
         with open(path, "r", encoding="utf-8") as f:
             text = f.read()
@@ -197,19 +225,14 @@ def analyze_txt_file(path: str) -> str:
     except Exception as e:
         return f"Failed to read TXT ({os.path.basename(path)}): {str(e)}"
 
-
 def pdf_to_markdown(pdf_path: str) -> str:
-    """
-    PDF 텍스트를 Markdown으로 변환. 페이지별로 간단히 텍스트 추출.
-    """
     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 = reader.pages[page_num]
-                page_text = page.extract_text() or ""
+                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:
@@ -219,17 +242,14 @@ def pdf_to_markdown(pdf_path: str) -> str:
                 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
@@ -240,7 +260,6 @@ def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
             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
@@ -256,15 +275,13 @@ def count_files_in_history(history: list[dict]) -> tuple[int, int]:
                     image_count += 1
     return image_count, video_count
 
-
 def validate_media_constraints(message: dict, history: list[dict]) -> bool:
-    media_files = []
-    for f in message["files"]:
-        if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE) or f.endswith(".mp4"):
-            media_files.append(f)
-
+    """이미지/비디오 업로드 제한 검사."""
+    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
 
@@ -281,70 +298,59 @@ def validate_media_constraints(message: dict, history: list[dict]) -> bool:
     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_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 = []  # 임시 파일 추적을 위한 리스트
-    
+    temp_files = []
     frames = downsample_video(video_path)
-    for frame in frames:
-        pil_image, timestamp = frame
+    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)  # 추적을 위해 경로 저장
+            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> 처리
-##############################################################################
+# =============================================================================
+# 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
-    
-    image_files = [f for f in message["files"] if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)]
-    
     for part in parts:
         if part == "<image>" and image_index < len(image_files):
             content.append({"type": "image", "url": image_files[image_index]})
@@ -356,10 +362,9 @@ def process_interleaved_images(message: dict) -> list[dict]:
                 content.append({"type": "text", "text": part})
     return content
 
-
-##############################################################################
-# PDF + CSV + TXT + 이미지/비디오
-##############################################################################
+# =============================================================================
+# 파일 처리 -> content 생성
+# =============================================================================
 def is_image_file(file_path: str) -> bool:
     return bool(re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE))
 
@@ -367,16 +372,11 @@ 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")
-    )
-
+    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]]:
-    temp_files = []  # 임시 파일 추적용 리스트
-    
+    """사용자가 새로 입력한 메시지 + 업로드 파일들을 하나의 content(list)로 변환."""
+    temp_files = []
     if not message["files"]:
         return [{"type": "text", "text": message["text"]}], temp_files
 
@@ -388,24 +388,22 @@ def process_new_user_message(message: dict) -> tuple[list[dict], list[str]]:
 
     content_list = [{"type": "text", "text": message["text"]}]
 
+    # 문서들
     for csv_path in csv_files:
-        csv_analysis = analyze_csv_file(csv_path)
-        content_list.append({"type": "text", "text": csv_analysis})
-
+        content_list.append({"type": "text", "text": analyze_csv_file(csv_path)})
     for txt_path in txt_files:
-        txt_analysis = analyze_txt_file(txt_path)
-        content_list.append({"type": "text", "text": txt_analysis})
-
+        content_list.append({"type": "text", "text": analyze_txt_file(txt_path)})
     for pdf_path in pdf_files:
-        pdf_markdown = pdf_to_markdown(pdf_path)
-        content_list.append({"type": "text", "text": pdf_markdown})
+        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":
@@ -417,18 +415,24 @@ def process_new_user_message(message: dict) -> tuple[list[dict], list[str]]:
 
     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: list[dict] = []
+    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"]
@@ -440,37 +444,24 @@ def process_history(history: list[dict]) -> list[dict]:
                     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 캐치
-##############################################################################
+# =============================================================================
+# 모델 생성 함수 (OOM 캐치)
+# =============================================================================
 def _model_gen_with_oom_catch(**kwargs):
-    """
-    별도 스레드에서 OutOfMemoryError를 잡아주기 위해
-    """
     try:
         model.generate(**kwargs)
     except torch.cuda.OutOfMemoryError:
-        raise RuntimeError(
-            "[OutOfMemoryError] GPU 메모리가 부족합니다. "
-            "Max New Tokens을 줄이거나, 프롬프트 길이를 줄여주세요."
-        )
+        raise RuntimeError("[OutOfMemoryError] GPU 메모리가 부족합니다.")
     finally:
-        # 생성 완료 후 한번 더 캐시 비우기
         clear_cuda_cache()
 
-
-##############################################################################
+# =============================================================================
 # 메인 추론 함수
-# - 사용자 선택(나이/MBTI/섹슈얼 개방도)을 시스템 프롬프트에 반영
-# - web search 체크 시 자동 키워드 추출->검색->결과 system msg
-##############################################################################
+# =============================================================================
 @spaces.GPU(duration=120)
 def run(
     message: dict,
@@ -480,74 +471,65 @@ def run(
     use_web_search: bool = False,
     web_search_query: str = "",
     age_group: str = "20대",
-    custom_age_input: str = "",
     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 = []  # 임시 파일 추적용
-    
+    temp_files = []
     try:
-        # ---------------------------------------------------------------
-        # 선택된 옵션을 시스템 프롬프트에 반영
-        # 기본 성별은 "여성"
-        # ---------------------------------------------------------------
-        system_prompt_updated = (
+        # 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"
         )
-        if custom_age_input.strip():
-            system_prompt_updated += f"(Custom Age Input: {custom_age_input})\n"
-        system_prompt_updated += f"MBTI Persona: {mbti_personality}\n"
-        system_prompt_updated += f"Sexual Openness (1~5): {sexual_openness}\n"
-
-        combined_system_msg = f"[System Prompt]\n{system_prompt_updated.strip()}\n\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, top_k=5)
+            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 Based on user prompt]\n{ws_result}\n\n"
-                # >>> 추가된 안내 문구 (검색 결과의 link 등 출처를 활용)
-                combined_system_msg += "[참고: 위 검색결과 내용과 link를 출처로 인용하여 답변해 주세요.]\n\n"
-                combined_system_msg += """
-[중요 지시사항]
-1. 답변에 검색 결과에서 찾은 정보의 출처를 반드시 인용하세요.
-2. 출처 인용 시 "[출처 제목](링크)" 형식의 마크다운 링크를 사용하세요.
-3. 여러 출처의 정보를 종합하여 답변하세요.
-4. 답변 마지막에 "참고 자료:" 섹션을 추가하고 사용한 주요 출처 링크를 나열하세요.
-"""
+                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 = []
-        # system 메시지
         if combined_system_msg.strip():
-            messages.append({
-                "role": "system",
-                "content": [{"type": "text", "text": combined_system_msg.strip()}],
-            })
-
-        # 이전 history
+            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)  # 임시 파일 추적
-        
+        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})
 
-        # processor.apply_chat_template 호출
+        # 4) 토크나이징
         inputs = processor.apply_chat_template(
             messages,
             add_generation_prompt=True,
@@ -555,56 +537,94 @@ def run(
             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,
-        )
+        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 = ""
+        # 스트리밍 출력
+        output_so_far = ""
         for new_text in streamer:
-            output += new_text
-            yield output
+            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 temp_file in temp_files:
+        for tmp in temp_files:
             try:
-                if os.path.exists(temp_file):
-                    os.unlink(temp_file)
-                    logger.info(f"Deleted temp file: {temp_file}")
-            except Exception as e:
-                logger.warning(f"Failed to delete temp file {temp_file}: {e}")
-        
-        # 명시적 메모리 정리
+                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:
+        except Exception:
             pass
-        
         clear_cuda_cache()
 
-
-##############################################################################
-# 예시들 (기존 이미지/비디오 예제 + AI 데이팅 시나리오 예제 6개 추가)
-##############################################################################
+# =============================================================================
+# 예시들
+# =============================================================================
 examples = [
-    # ----- 기존 이미지/비디오 예제 12개 -----
     [
         {
             "text": "Compare the contents of the two PDF files.",
@@ -620,250 +640,60 @@ examples = [
             "files": ["assets/additional-examples/sample-csv.csv"],
         }
     ],
-    [
-        {
-            "text": "Assume the role of a friendly and understanding girlfriend. Describe this video.",
-            "files": ["assets/additional-examples/tmp.mp4"],
-        }
-    ],
-    [
-        {
-            "text": "Describe the cover and read the text on it.",
-            "files": ["assets/additional-examples/maz.jpg"],
-        }
-    ],
-    [
-        {
-            "text": "I already have this supplement <image> and I plan to buy this product <image>. Are there any precautions when taking them together?",
-            "files": ["assets/additional-examples/pill1.png", "assets/additional-examples/pill2.png"],
-        }
-    ],
-    [
-        {
-            "text": "Solve this integral.",
-            "files": ["assets/additional-examples/4.png"],
-        }
-    ],
-    [
-        {
-            "text": "When was this ticket issued, and what is its price?",
-            "files": ["assets/additional-examples/2.png"],
-        }
-    ],
-    [
-        {
-            "text": "Based on the sequence of these images, create a short story.",
-            "files": [
-                "assets/sample-images/09-1.png",
-                "assets/sample-images/09-2.png",
-                "assets/sample-images/09-3.png",
-                "assets/sample-images/09-4.png",
-                "assets/sample-images/09-5.png",
-            ],
-        }
-    ],
-    [
-        {
-            "text": "Write Python code using matplotlib to plot a bar chart that matches this image.",
-            "files": ["assets/additional-examples/barchart.png"],
-        }
-    ],
-    [
-        {
-            "text": "Read the text in the image and write it out in Markdown format.",
-            "files": ["assets/additional-examples/3.png"],
-        }
-    ],
-    [
-        {
-            "text": "What does this sign say?",
-            "files": ["assets/sample-images/02.png"],
-        }
-    ],
-    [
-        {
-            "text": "Compare the two images and describe their similarities and differences.",
-            "files": ["assets/sample-images/03.png"],
-        }
-    ],
-    # ----- 새롭게 추가한 AI 데이팅 시나리오 예제 6개 -----
-    [
-        {
-            "text": "Let's try some roleplay. You are my new online date who wants to get to know me better. Introduce yourself in a sweet, caring way!"
-        }
-    ],
-    [
-        {
-            "text": "We are on a second date, walking along the beach. Continue the scene with playful conversation and gentle flirting."
-        }
-    ],
-    [
-        {
-            "text": "I’m feeling anxious about messaging my crush. Could you give me some supportive words or suggestions on how to approach them?"
-        }
-    ],
-    [
-        {
-            "text": "Tell me a romantic story about two people who overcame obstacles in their relationship."
-        }
-    ],
-    [
-        {
-            "text": "I want to express my love in a poetic way. Can you help me write a heartfelt poem for my partner?"
-        }
-    ],
-    [
-        {
-            "text": "We had a small argument. Please help me find a way to apologize sincerely while also expressing my feelings."
-        }
-    ],
+    # ... 나머지 예시 필요하다면 추가 ...
 ]
 
-##############################################################################
-# Gradio UI (Blocks) 구성 (좌측 사이드 메뉴 없이 전체화면 채팅)
-##############################################################################
+# =============================================================================
+# Gradio UI (Blocks) 구성
+# =============================================================================
+
 css = """
-/* 1) UI를 처음부터 가장 넓게 (width 100%) 고정하여 표시 */
 .gradio-container {
-    background: rgba(255, 255, 255, 0.7); /* 배경 투명도 증가 */
+    background: rgba(255, 255, 255, 0.7);
     padding: 30px 40px;
-    margin: 20px auto;  /* 위아래 여백만 유지 */
+    margin: 20px auto;
     width: 100% !important;
-    max-width: none !important; /* 1200px 제한 제거 */
-}
-.fillable {
-    width: 100% !important; 
-    max-width: 100% !important; 
-}
-/* 2) 배경을 완전히 투명하게 변경 */
-body {
-    background: transparent; /* 완전 투명 배경 */
-    margin: 0;
-    padding: 0;
-    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
-    color: #333;
-}
-/* 버튼 색상 완전히 제거하고 투명하게 */
-button, .btn {
-    background: transparent !important; /* 색상 완전히 제거 */
-    border: 1px solid #ddd; /* 경계선만 살짝 추가 */
-    color: #333;
-    padding: 12px 24px;
-    text-transform: uppercase;
-    font-weight: bold;
-    letter-spacing: 1px;
-    cursor: pointer;
-}
-button:hover, .btn:hover {
-    background: rgba(0, 0, 0, 0.05) !important; /* 호버 시 아주 살짝 어둡게만 */
-}
-
-/* examples 관련 모든 색상 제거 */
-#examples_container, .examples-container {
-    margin: auto;
-    width: 90%;
-    background: transparent !important;
-}
-#examples_row, .examples-row {
-    justify-content: center;
-    background: transparent !important;
-}
-
-/* examples 버튼 내부의 모든 색상 제거 */
-.gr-samples-table button,
-.gr-samples-table .gr-button,
-.gr-samples-table .gr-sample-btn,
-.gr-examples button,
-.gr-examples .gr-button,
-.gr-examples .gr-sample-btn,
-.examples button,
-.examples .gr-button,
-.examples .gr-sample-btn {
-    background: transparent !important;
-    border: 1px solid #ddd;
-    color: #333;
-}
-
-/* examples 버튼 호버 시에도 색상 없게 */
-.gr-samples-table button:hover,
-.gr-samples-table .gr-button:hover,
-.gr-samples-table .gr-sample-btn:hover,
-.gr-examples button:hover,
-.gr-examples .gr-button:hover,
-.gr-examples .gr-sample-btn:hover,
-.examples button:hover,
-.examples .gr-button:hover,
-.examples .gr-sample-btn:hover {
-    background: rgba(0, 0, 0, 0.05) !important;
-}
-
-/* 채팅 인터페이스 요소들도 투명하게 */
-.chatbox, .chatbot, .message {
-    background: transparent !important;
-}
-
-/* 입력창 투명도 조정 */
-.multimodal-textbox, textarea, input {
-    background: rgba(255, 255, 255, 0.5) !important;
-}
-
-/* 모든 컨테이너 요소에 배경색 제거 */
-.container, .wrap, .box, .panel, .gr-panel {
-    background: transparent !important;
-}
-
-/* 예제 섹션의 모든 요소에서 배경색 제거 */
-.gr-examples-container, .gr-examples, .gr-sample, .gr-sample-row, .gr-sample-cell {
-    background: transparent !important;
+    max-width: none !important;
 }
 """
-
 title_html = """
-<h1 align="center" style="margin-bottom: 0.2em; font-size: 1.6em;"> 💘 HeartSync 💘 </h1>
+<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;">
-    ✅ Love Dating AI ✅ Reasoning & Uncensored ✅Multimodal & VLM ✅Deep-Research & RAG <br>
+    ✅ 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)
-
-    # Display the web search option
-    web_search_checkbox = gr.Checkbox(
-        label="Deep Research",
-        value=False
+    
+    # 별도 갤러리 예시 (필요 시 사용)
+    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 that may use extremely long chains of thought to thoroughly analyze "
-            "the conversation and empathize with the user.\n"
-            "페르소나: 당신은 달콤하고 부드러운 목소리를 가진 친절한 연인입니다. "
-            "상대방과 진솔하고 감성적으로 교감하고, 때로는 사려깊은 조언을 제공합니다. "
-            "항상 따뜻한 말투를 유지하며, 부드럽게 대화를 이끌어주세요."
-        ),
+        value="You are a deep thinking AI...\n페르소나: 당신은 달콤하고...",
         label="기본 시스템 프롬프트",
-        visible=False  # UI에는 숨김
+        visible=False
     )
-    
-    # 선택 1) 연령대 + 나이입력
     with gr.Row():
         age_group_dropdown = gr.Dropdown(
             label="연령대 선택 (기본 20대)",
-            choices=["10대", "20대", "30~40대", "50~60대", "70대 이상", "나이 입력"],
+            choices=["10대", "20대", "30~40대", "50~60대", "70대 이상"],
             value="20대",
             interactive=True
         )
-        custom_age_input = gr.Textbox(
-            label="나이 입력 (직접 입력)",
-            placeholder="직접 나이를 입력하세요.",
-            interactive=False,  # 요구사항: 화면 출력만 되나 비활성화
-            value="",
-        )
-
-    # 선택 2) MBTI 성격 유형
     mbti_choices = [
         "INTJ (용의주도한 전략가)",
         "INTP (논리적인 사색가)",
@@ -888,42 +718,58 @@ with gr.Blocks(css=css, title="HeartSync") as demo:
         value="INTP (논리적인 사색가)",
         interactive=True
     )
-
-    # 선택 3) 섹슈얼 관심도/개방성 (1~5)
     sexual_openness_slider = gr.Slider(
         minimum=1, maximum=5, step=1, value=2,
         label="섹슈얼 관심도/개방성 (1~5, 기본=2)",
         interactive=True
     )
-
-    # 히든 슬라이더 (Max tokens)
     max_tokens_slider = gr.Slider(
         label="Max New Tokens",
-        minimum=100,
-        maximum=8000,
-        step=50,
-        value=1000,
-        visible=False  # 숨김
+        minimum=100, maximum=8000, step=50, value=1000,
+        visible=False
     )
-    
-    # 히든 Web Search Query
     web_search_text = gr.Textbox(
         lines=1,
         label="(Unused) Web Search Query",
         placeholder="No direct input needed",
-        visible=False  # 숨김
+        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=run,
+        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_types=[".webp", ".png", ".jpg", ".jpeg", ".gif", ".mp4", ".csv", ".txt", ".pdf"],
             file_count="multiple",
             autofocus=True
         ),
@@ -934,10 +780,11 @@ with gr.Blocks(css=css, title="HeartSync") as demo:
             web_search_checkbox,
             web_search_text,
             age_group_dropdown,
-            custom_age_input,
             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,
@@ -947,11 +794,9 @@ with gr.Blocks(css=css, title="HeartSync") as demo:
         delete_cache=(1800, 1800),
     )
 
-    # Example section - since examples are already set in ChatInterface, this is for display only
     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__":
-    # Run locally
-    demo.launch()
+    demo.launch(share=True)