app.py

import gradio as gr
import replicate
from PIL import Image
from google import genai
from google.genai import types
import io
import base64
import tempfile
import os
import uuid
import requests
from io import BytesIO
import time
import json
import datetime

# 환경변수에서 API 토큰과 비밀번호 가져오기
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
PASSWORD = os.getenv("APP_PASSWORD")  # 기본값 설정

# API 키 검증
def validate_api_keys():
    """API 키들이 올바르게 설정되었는지 확인"""
    missing_keys = []
    
    if not REPLICATE_API_TOKEN:
        missing_keys.append("REPLICATE_API_TOKEN")
    
    if not GEMINI_API_KEY:
        missing_keys.append("GEMINI_API_KEY")
    
    if missing_keys:
        error_msg = f"다음 환경변수가 설정되지 않았습니다: {', '.join(missing_keys)}"
        raise ValueError(error_msg)
    
    return True

def log_message(message, log_list=None):
    """로그 메시지를 타임스탬프와 함께 기록"""
    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
    formatted_message = f"[{timestamp}] {message}"
    print(formatted_message)  # 콘솔에도 출력
    if log_list is not None:
        log_list.append(formatted_message)
    return formatted_message

def translate_to_english(korean_prompt, log_list):
    """한국어 프롬프트를 FLUX 이미지 생성에 최적화된 영어로 번역"""
    log_message("=== 번역 프로세스 시작 ===", log_list)
    log_message(f"입력된 한국어 프롬프트: '{korean_prompt}'", log_list)
    
    try:
        log_message("Gemini API 클라이언트 초기화 시작...", log_list)
        client = genai.Client(api_key=GEMINI_API_KEY)
        log_message("Gemini API 클라이언트 초기화 완료", log_list)
        
        system_instruction = """You are a professional prompt translator for FLUX image generation models. 
        Translate Korean image editing prompts to optimized English prompts that work best with FLUX models.
        Keep the translation:
        - Clear and specific
        - Descriptive but concise
        - Focused on visual elements
        - Suitable for AI image generation
        
        Only return the translated English prompt, nothing else."""
        
        log_message("번역 요청 준비 완료", log_list)
        log_message("Gemini API 호출 시작...", log_list)
        
        start_time = time.time()
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            config=types.GenerateContentConfig(
                system_instruction=system_instruction,
                temperature=0.3
            ),
            contents=[f"Translate this Korean prompt to English for FLUX image editing: {korean_prompt}"]
        )
        end_time = time.time()
        
        log_message(f"Gemini API 응답 수신 완료 (소요시간: {end_time - start_time:.2f}초)", log_list)
        
        translated_text = response.text.strip()
        log_message(f"번역 결과: '{translated_text}'", log_list)
        log_message("=== 번역 프로세스 완료 ===", log_list)
        
        return translated_text
        
    except Exception as e:
        error_msg = f"번역 중 오류 발생: {str(e)}"
        log_message(error_msg, log_list)
        log_message("원본 한국어 프롬프트를 그대로 사용합니다", log_list)
        log_message("=== 번역 프로세스 실패 ===", log_list)
        return korean_prompt

def upscale_image(image_path, output_format, english_prompt, log_list):
    """Clarity Upscaler를 사용하여 이미지 업스케일링 (모든 매개변수 고정)"""
    log_message("=== 화질 개선(업스케일링) 프로세스 시작 ===", log_list)
    log_message(f"입력 이미지 경로: {image_path}", log_list)
    log_message(f"출력 포맷: {output_format}", log_list)
    log_message(f"사용할 프롬프트: '{english_prompt}'", log_list)
    
    try:
        # 파일 존재 확인
        if not os.path.exists(image_path):
            error_msg = f"입력 이미지 파일이 존재하지 않습니다: {image_path}"
            log_message(error_msg, log_list)
            return None
            
        # 파일 크기 확인
        file_size = os.path.getsize(image_path)
        log_message(f"입력 파일 크기: {file_size} bytes ({file_size/1024/1024:.2f} MB)", log_list)
        
        log_message("Replicate API 클라이언트 초기화 시작...", log_list)
        client = replicate.Client(api_token=REPLICATE_API_TOKEN)
        log_message("Replicate API 클라이언트 초기화 완료", log_list)
        
        log_message("이미지 파일 읽기 시작...", log_list)
        with open(image_path, "rb") as file:
            log_message("이미지 파일 읽기 완료", log_list)
            
            # 업스케일 파라미터 설정
            input_data = {
                "image": file,
                "scale_factor": 2,  # 고정: 2배 확대
                "resemblance": 0.8,  # 고정: 높은 원본 유사도
                "creativity": 0.2,   # 고정: 낮은 창의성
                "output_format": output_format.lower(),
                "prompt": english_prompt,
                "negative_prompt": "(worst quality, low quality, normal quality:2)"
            }
            
            log_message("업스케일 파라미터 설정:", log_list)
            log_message(f"  - scale_factor: 2", log_list)
            log_message(f"  - resemblance: 0.8", log_list)
            log_message(f"  - creativity: 0.2", log_list)
            log_message(f"  - output_format: {output_format.lower()}", log_list)
            log_message(f"  - negative_prompt: (worst quality, low quality, normal quality:2)", log_list)
            
            log_message("Clarity Upscaler API 호출 시작...", log_list)
            log_message("모델: philz1337x/clarity-upscaler", log_list)
            
            start_time = time.time()
            try:
                output = client.run(
                    "philz1337x/clarity-upscaler:dfad41707589d68ecdccd1dfa600d55a208f9310748e44bfe35b4a6291453d5e",
                    input=input_data
                )
                end_time = time.time()
                log_message(f"Clarity Upscaler API 응답 수신 완료 (소요시간: {end_time - start_time:.2f}초)", log_list)
                
            except Exception as api_error:
                error_msg = f"Clarity Upscaler API 호출 실패: {str(api_error)}"
                log_message(error_msg, log_list)
                log_message("=== 화질 개선 프로세스 실패 ===", log_list)
                return None
        
        # API 응답 검증
        log_message(f"API 응답 타입: {type(output)}", log_list)
        log_message(f"API 응답 내용: {output}", log_list)
        
        # 결과 이미지 다운로드
        if output and isinstance(output, list) and len(output) > 0:
            result_url = output[0]
            log_message(f"결과 이미지 URL: {result_url}", log_list)
            log_message("결과 이미지 다운로드 시작...", log_list)
            
            try:
                download_start = time.time()
                response = requests.get(result_url, timeout=30)
                download_end = time.time()
                
                log_message(f"다운로드 응답 상태: {response.status_code}", log_list)
                log_message(f"다운로드 소요시간: {download_end - download_start:.2f}초", log_list)
                log_message(f"다운로드된 데이터 크기: {len(response.content)} bytes ({len(response.content)/1024/1024:.2f} MB)", log_list)
                
                if response.status_code == 200:
                    log_message("이미지 데이터를 PIL Image로 변환 중...", log_list)
                    result_image = Image.open(BytesIO(response.content))
                    log_message(f"변환된 이미지 크기: {result_image.size}", log_list)
                    log_message(f"변환된 이미지 모드: {result_image.mode}", log_list)
                    
                    # 임시 파일로 저장
                    ext = output_format.lower()
                    upscaled_filename = f"upscaled_temp_{uuid.uuid4()}.{ext}"
                    log_message(f"업스케일된 이미지 저장 파일명: {upscaled_filename}", log_list)
                    
                    # RGBA 모드 처리
                    if ext == 'jpg' and result_image.mode == 'RGBA':
                        log_message("RGBA 모드를 RGB로 변환 중 (JPG 저장을 위해)...", log_list)
                        result_image = result_image.convert('RGB')
                    
                    # 이미지 저장
                    save_start = time.time()
                    if ext == 'jpg':
                        result_image.save(upscaled_filename, format='JPEG', quality=95)
                    else:
                        result_image.save(upscaled_filename, format='PNG')
                    save_end = time.time()
                    
                    log_message(f"이미지 저장 완료 (소요시간: {save_end - save_start:.2f}초)", log_list)
                    
                    # 저장된 파일 크기 확인
                    saved_size = os.path.getsize(upscaled_filename)
                    log_message(f"저장된 파일 크기: {saved_size} bytes ({saved_size/1024/1024:.2f} MB)", log_list)
                    
                    log_message("=== 화질 개선 프로세스 완료 ===", log_list)
                    return upscaled_filename
                else:
                    error_msg = f"이미지 다운로드 실패 - HTTP 상태: {response.status_code}"
                    log_message(error_msg, log_list)
                    log_message("=== 화질 개선 프로세스 실패 ===", log_list)
                    return None
                    
            except requests.exceptions.Timeout:
                log_message("이미지 다운로드 타임아웃 (30초)", log_list)
                log_message("=== 화질 개선 프로세스 실패 ===", log_list)
                return None
            except Exception as download_error:
                error_msg = f"이미지 다운로드 중 오류: {str(download_error)}"
                log_message(error_msg, log_list)
                log_message("=== 화질 개선 프로세스 실패 ===", log_list)
                return None
        else:
            log_message("API 응답이 비어있거나 예상 형식과 다릅니다", log_list)
            if output:
                log_message(f"실제 응답: {json.dumps(output, indent=2) if isinstance(output, dict) else str(output)}", log_list)
            log_message("=== 화질 개선 프로세스 실패 ===", log_list)
            return None
        
    except Exception as e:
        error_msg = f"업스케일링 프로세스 중 예상치 못한 오류: {str(e)}"
        log_message(error_msg, log_list)
        log_message("=== 화질 개선 프로세스 실패 ===", log_list)
        return None

def edit_image(input_image, password, korean_prompt, output_format, aspect_ratio, upscale_option, current_images, current_downloads):
    log_list = []
    log_message("=" * 60, log_list)
    log_message("새로운 이미지 편집 작업 시작", log_list)
    log_message("=" * 60, log_list)
    
    # API 키 검증
    try:
        validate_api_keys()
        log_message("API 키 검증 완료", log_list)
    except ValueError as e:
        error_msg = str(e)
        log_message(f"API 키 검증 실패: {error_msg}", log_list)
        return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
    
    # 입력 검증
    log_message("=== 입력 검증 단계 ===", log_list)
    
    if not password:
        error_msg = "비밀번호를 입력하세요."
        log_message(f"검증 실패: {error_msg}", log_list)
        return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
    
    if password != PASSWORD:
        error_msg = "비밀번호가 올바르지 않습니다."
        log_message(f"검증 실패: {error_msg}", log_list)
        return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
    
    log_message("비밀번호 검증 통과", log_list)
    
    if input_image is None:
        error_msg = "이미지를 업로드하세요."
        log_message(f"검증 실패: {error_msg}", log_list)
        return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
    
    log_message(f"입력 이미지 경로: {input_image}", log_list)
    
    if not korean_prompt.strip():
        error_msg = "편집 지시사항을 입력하세요."
        log_message(f"검증 실패: {error_msg}", log_list)
        return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
    
    log_message(f"편집 지시사항: '{korean_prompt.strip()}'", log_list)
    log_message(f"출력 포맷: {output_format}", log_list)
    log_message(f"화면 비율: {aspect_ratio}", log_list)
    log_message(f"화질 개선 옵션: {upscale_option}", log_list)
    log_message("모든 입력 검증 통과", log_list)
    
    try:
        # Replicate 클라이언트 설정
        log_message("=== Replicate 설정 단계 ===", log_list)
        log_message("Replicate 클라이언트 초기화 시작...", log_list)
        client = replicate.Client(api_token=REPLICATE_API_TOKEN)
        log_message("Replicate 클라이언트 초기화 완료", log_list)
        
        # 한국어 프롬프트를 영어로 번역
        english_prompt = translate_to_english(korean_prompt, log_list)
        
        # 입력 이미지 정보 확인
        log_message("=== 입력 이미지 분석 ===", log_list)
        try:
            with Image.open(input_image) as img:
                log_message(f"이미지 크기: {img.size}", log_list)
                log_message(f"이미지 모드: {img.mode}", log_list)
                log_message(f"이미지 포맷: {img.format}", log_list)
        except Exception as img_error:
            log_message(f"이미지 정보 읽기 오류: {str(img_error)}", log_list)
        
        # 파일 크기 확인
        file_size = os.path.getsize(input_image)
        log_message(f"입력 파일 크기: {file_size} bytes ({file_size/1024/1024:.2f} MB)", log_list)
        
        # FLUX 이미지 편집 시작
        log_message("=== FLUX 이미지 편집 단계 ===", log_list)
        log_message("이미지 파일 읽기 시작...", log_list)
        
        with open(input_image, "rb") as file:
            input_file = file
            log_message("파일 읽기 완료", log_list)
            
            # Replicate API 호출 파라미터 설정
            input_data = {
                "prompt": english_prompt,
                "input_image": input_file,
                "output_format": output_format.lower(),
                "aspect_ratio": aspect_ratio
            }
            
            
            flux_start_time = time.time()
            try:
                output = client.run(
                    "black-forest-labs/flux-kontext-pro",
                    input=input_data
                )
                flux_end_time = time.time()
                log_message(f"FLUX API 응답 수신 완료 (소요시간: {flux_end_time - flux_start_time:.2f}초)", log_list)
                
            except Exception as flux_error:
                error_msg = f"FLUX API 호출 실패: {str(flux_error)}"
                log_message(error_msg, log_list)
                return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
        
        log_message(f"FLUX API 응답 타입: {type(output)}", log_list)
        
        # 결과 이미지 처리
        log_message("=== 결과 이미지 처리 단계 ===", log_list)
        temp_filename = f"temp_output_{uuid.uuid4()}.{output_format.lower()}"
        log_message(f"임시 파일명: {temp_filename}", log_list)
        
        try:
            log_message("결과 이미지 저장 시작...", log_list)
            with open(temp_filename, "wb") as file:
                file.write(output.read())
            
            # 저장된 파일 크기 확인
            saved_size = os.path.getsize(temp_filename)
            log_message(f"FLUX 결과 이미지 저장 완료", log_list)
            log_message(f"저장된 파일 크기: {saved_size} bytes ({saved_size/1024/1024:.2f} MB)", log_list)
            
        except Exception as save_error:
            error_msg = f"결과 이미지 저장 실패: {str(save_error)}"
            log_message(error_msg, log_list)
            return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
        
        # 결과 이미지 로드
        try:
            log_message("결과 이미지를 PIL Image로 로드 중...", log_list)
            result_image = Image.open(temp_filename)
            log_message(f"로드된 이미지 크기: {result_image.size}", log_list)
            log_message(f"로드된 이미지 모드: {result_image.mode}", log_list)
            
        except Exception as load_error:
            error_msg = f"결과 이미지 로드 실패: {str(load_error)}"
            log_message(error_msg, log_list)
            return None, None, "\n".join(log_list) + f"\n오류: {error_msg}", current_images, current_downloads
        
        # 업스케일링 적용 (선택사항)
        final_image = result_image
        if upscale_option == "적용":
            upscaled_path = upscale_image(temp_filename, output_format, english_prompt, log_list)
            
            if upscaled_path:
                try:
                    log_message("업스케일된 이미지 로드 중...", log_list)
                    final_image = Image.open(upscaled_path)
                    log_message(f"최종 이미지 크기: {final_image.size}", log_list)
                    log_message("화질 개선이 성공적으로 적용되었습니다", log_list)
                    
                    # 업스케일 임시 파일 정리
                    try:
                        os.remove(upscaled_path)
                        log_message("업스케일 임시 파일 정리 완료", log_list)
                    except:
                        pass
                        
                except Exception as upscale_load_error:
                    log_message(f"업스케일된 이미지 로드 실패: {str(upscale_load_error)}", log_list)
                    log_message("원본 FLUX 결과 이미지를 사용합니다", log_list)
            else:
                log_message("화질 개선 실패 - 원본 FLUX 결과 이미지를 사용합니다", log_list)
        else:
            log_message("화질 개선 옵션이 선택되지 않음 - 원본 FLUX 결과 사용", log_list)
        
        # 자동 저장 처리
        log_message("=== 결과 저장 단계 ===", log_list)
        
        if current_images is None:
            current_images = []
        if current_downloads is None:
            current_downloads = []
        
        # 갤러리용 임시 파일 생성
        gallery_temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
        final_image.save(gallery_temp_file.name, 'png')
        current_images.append(gallery_temp_file.name)
        log_message(f"갤러리용 이미지 저장: {gallery_temp_file.name}", log_list)
        
        # 다운로드용 파일 생성
        download_filename = f"edited_image_{len(current_downloads) + 1}.{output_format.lower()}"
        final_image.save(download_filename)
        current_downloads.append(download_filename)
        log_message(f"다운로드용 이미지 저장: {download_filename}", log_list)
        
        # 임시 파일 정리
        try:
            os.remove(temp_filename)
            log_message("FLUX 임시 파일 정리 완료", log_list)
        except:
            pass
        
        log_message("=" * 60, log_list)
        log_message("이미지 편집 작업 성공적으로 완료!", log_list)
        log_message("=" * 60, log_list)
        
        return final_image, current_downloads, "\n".join(log_list), current_images, current_downloads
        
    except Exception as e:
        error_msg = f"예상치 못한 오류 발생: {str(e)}"
        log_message(error_msg, log_list)
        log_message("=" * 60, log_list)
        log_message("이미지 편집 작업 실패", log_list)
        log_message("=" * 60, log_list)
        return None, current_downloads, "\n".join(log_list), current_images, current_downloads

def save_output_as_input(output_image, current_images, current_downloads):
    """출력 이미지를 입력으로 저장"""
    log_list = []
    log_message("=== 출력 이미지를 입력으로 저장 시작 ===", log_list)
    
    if output_image is None:
        log_message("저장할 출력 이미지가 없습니다", log_list)
        return None, current_images, current_downloads
    
    try:
        log_message(f"출력 이미지 타입: {type(output_image)}", log_list)
        
        # numpy array를 PIL Image로 변환
        if hasattr(output_image, 'shape'):  # numpy array인 경우
            log_message("numpy array를 PIL Image로 변환 중...", log_list)
            output_image = Image.fromarray(output_image)
        
        log_message(f"변환된 이미지 크기: {output_image.size}", log_list)
        log_message(f"변환된 이미지 모드: {output_image.mode}", log_list)
        
        # 임시 파일로 저장 (입력용)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
        output_image.save(temp_file.name, 'png')
        
        file_size = os.path.getsize(temp_file.name)
        log_message(f"입력용 이미지 저장 완료: {temp_file.name}", log_list)
        log_message(f"저장된 파일 크기: {file_size} bytes ({file_size/1024/1024:.2f} MB)", log_list)
        log_message("=== 저장 프로세스 완료 ===", log_list)
        
        print("\n".join(log_list))  # 콘솔에 로그 출력
        return temp_file.name, current_images, current_downloads
        
    except Exception as e:
        error_msg = f"이미지 저장 중 오류: {str(e)}"
        log_message(error_msg, log_list)
        log_message("=== 저장 프로세스 실패 ===", log_list)
        print("\n".join(log_list))  # 콘솔에 로그 출력
        return None, current_images, current_downloads

# 애플리케이션 시작 시 API 키 확인
try:
    validate_api_keys()
    print("✅ 모든 API 키가 올바르게 설정되었습니다.")
except ValueError as e:
    print(f"❌ {e}")

# Gradio 인터페이스
with gr.Blocks(title="이미지 편집기 (보안 강화 버전)") as demo:
    # 상태 변수들
    saved_images = gr.State([])
    saved_downloads = gr.State([])
    
    gr.Markdown("# 🎨 AI 이미지 편집기 (보안 강화 버전)")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="📤 입력 이미지")
            password = gr.Textbox(
                label="🔐 비밀번호", 
                type="password",
                placeholder="비밀번호를 입력하세요"
            )
            korean_prompt = gr.Textbox(
                label="✏️ 편집 지시사항 (한국어)",
                placeholder="예: 이 사람을 만화 캐릭터로 바꿔줘",
                lines=3
            )
            
            with gr.Row():
                output_format = gr.Radio(
                    choices=["jpg", "png"],
                    value="jpg",
                    label="📁 출력 포맷"
                )
                aspect_ratio = gr.Dropdown(
                    choices=["match_input_image", "1:1", "3:2", "2:3"],
                    value="match_input_image",
                    label="📐 화면 비율"
                )
            
            upscale_option = gr.Radio(
                choices=["없음", "적용"],
                value="없음",
                label="🔍 화질 개선 (2배 확대)"
            )
            
            edit_btn = gr.Button("🚀 이미지 편집", variant="primary", size="lg")
        
        with gr.Column():
            output_image = gr.Image(label="✨ 편집된 이미지")
            log_output = gr.Textbox(
                label="📋 상세 로그", 
                lines=15, 
                max_lines=20,
                show_copy_button=True
            )
            download_files = gr.File(label="💾 다운로드", file_count="multiple")
            
            save_btn = gr.Button("🔄 출력 이미지를 입력으로 저장", variant="secondary")
    
    # 저장된 이미지들을 표시하는 갤러리
    with gr.Row():
        output_gallery = gr.Gallery(
            label="📸 편집 기록", 
            show_label=True, 
            elem_id="gallery",
            columns=3,
            rows=2,
            height="auto"
        )
    
    # 이벤트 핸들러
    edit_btn.click(
        fn=edit_image,
        inputs=[input_image, password, korean_prompt, output_format, aspect_ratio, upscale_option, saved_images, saved_downloads],
        outputs=[output_image, download_files, log_output, saved_images, saved_downloads]
    )
    
    save_btn.click(
        fn=save_output_as_input,
        inputs=[output_image, saved_images, saved_downloads],
        outputs=[input_image, saved_images, saved_downloads]
    )
    
    # 갤러리 업데이트
    saved_images.change(
        fn=lambda images: images if images else [],
        inputs=[saved_images],
        outputs=[output_gallery]
    )

if __name__ == "__main__":
    print("🚀 이미지 편집 애플리케이션 시작 중...")

    # API 키 상태 확인
    try:
        validate_api_keys()
        print("✅ 모든 API 키가 설정되어 애플리케이션을 시작합니다.")
    except ValueError as e:
        print(f"⚠️  경고: {e}")
        print("일부 기능이 제한될 수 있습니다.")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )