app.py

import gradio as gr
import replicate
import os
import requests
from PIL import Image
import io

# Replicate API 설정
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
if not REPLICATE_API_TOKEN:
    raise ValueError("REPLICATE_API_TOKEN이 설정되지 않았습니다. Spaces Settings에서 Secret을 추가해주세요.")
client = replicate.Client()

def upload_file_to_replicate(file_path):
    """파일을 임시 URL로 변환 (base64 data URL 사용)"""
    if file_path is None:
        return None
    
    try:
        import base64
        with open(file_path, "rb") as file:
            file_data = file.read()
            file_base64 = base64.b64encode(file_data).decode()
            # 파일 확장자 확인
            if file_path.lower().endswith('.png'):
                mime_type = "image/png"
            elif file_path.lower().endswith('.jpg') or file_path.lower().endswith('.jpeg'):
                mime_type = "image/jpeg"
            elif file_path.lower().endswith('.webp'):
                mime_type = "image/webp"
            elif file_path.lower().endswith('.gif'):
                mime_type = "image/gif"
            else:
                mime_type = "image/png"  # 기본값
            
            data_url = f"data:{mime_type};base64,{file_base64}"
            return data_url
    except Exception as e:
        print(f"파일 변환 오류: {e}")
        return None

def generate_pattern(prompt, reference_image=None):
    """보도블록 패턴 생성 (2장)"""
    try:
        # 보도블록 패턴 생성을 위한 프롬프트 강화
        enhanced_prompt = f"sidewalk block pattern, paving stone design, {prompt}, top-down view, seamless tile pattern, architectural blueprint style"
        
        results = []
        
        # 참고 이미지가 있을 경우 URL로 변환
        reference_url = None
        if reference_image is not None:
            reference_url = upload_file_to_replicate(reference_image)
            if reference_url is None:
                return None, None, "참고 이미지 업로드에 실패했습니다."
        
        for i in range(2):  # 2장 생성
            if reference_url is not None:
                # 참고 이미지가 있을 경우
                output = client.run(
                    "black-forest-labs/flux-kontext-dev",
                    input={
                        "prompt": f"{enhanced_prompt}, variation {i+1}",
                        "input_image": reference_url,
                        "aspect_ratio": "1:1",
                        "num_inference_steps": 25,
                        "guidance": 3.0,
                        "output_format": "png",
                        "go_fast": True
                    }
                )
                # FileOutput을 URL 문자열로 변환
                result_url = str(output)
            else:
                # 텍스트만으로 생성 (SDXL 사용)
                output = client.run(
                    "stability-ai/sdxl:7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc",
                    input={
                        "prompt": f"{enhanced_prompt}, variation {i+1}",
                        "negative_prompt": "blurry, low quality, distorted, people, cars, buildings",
                        "width": 512,
                        "height": 512,
                        "num_inference_steps": 20,
                        "guidance_scale": 7.5
                    }
                )
                # SDXL은 리스트로 리턴되므로 첫 번째 요소 선택
                if isinstance(output, list):
                    result_url = str(output[0])
                else:
                    result_url = str(output)
            
            results.append(result_url)
        
        return results[0], results[1], "패턴 생성 완료! 마음에 드는 패턴을 다운로드하세요."
        
    except Exception as e:
        return None, None, f"오류가 발생했습니다: {str(e)}"

def apply_pattern_to_road(road_image, pattern_image):
    """도로에 패턴 적용"""
    try:
        # 이미지들을 URL로 변환
        road_url = upload_file_to_replicate(road_image)
        pattern_url = upload_file_to_replicate(pattern_image)
        
        if road_url is None:
            return None, "도로 이미지 업로드에 실패했습니다."
        if pattern_url is None:
            return None, "패턴 이미지 업로드에 실패했습니다."
        
        # 예시 스타일로 명확한 프롬프트
        # preset_prompt = "Transform the whole sidewalk pavement in the first image to have the exact same decorative pattern as shown in the second image. Replace the whole existing pavement with the new pattern. Combine these photos into one fluid scene."
        preset_prompt = "Combine these photos into one fluid scene. Transform the sidewalk pavement in the first image to have the exact same sidewalk block pattern as shown in the second image. Replace the existing pavement with the new sidewalk block pattern."
        # Multi-image list FLUX Kontext로 패턴 적용
        output = client.run(
            "flux-kontext-apps/multi-image-list",
            input={
                "prompt": preset_prompt,
                "input_images": [road_url, pattern_url],  # 리스트로 전달
                "aspect_ratio": "4:3",
                "output_format": "png",
                "safety_tolerance": 2
            }
        )
        
        # FileOutput을 URL 문자열로 변환
        result_url = str(output)
        
        return result_url, "패턴 적용 완료! Before/After 비교해보세요."
        
    except Exception as e:
        return None, f"오류가 발생했습니다: {str(e)}"

# CSS 스타일링
css = """
#col-container {
    margin: 0 auto;
    max-width: 900px;
}
.pattern-gallery {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 10px;
}
.header-text {
    text-align: center;
    margin-bottom: 20px;
    padding: 20px;
    border-radius: 15px;
    border: 2px solid #e0e0e0;
}
.logo-container {
    display: flex;
    align-items: center;
    justify-content: center;
    gap: 15px;
    margin-bottom: 15px;
}
.logo-img {
    width: 60px;
    height: 60px;
    border-radius: 10px;
}
"""

with gr.Blocks(css=css, title="GoBlocks - 보도블록 리디자인 서비스") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        <h2 style="margin: 10px 0;">보도블럭 조감도 & 상세도 제작 서비스</h2>
        <p style="margin: 5px 0; font-size: 1.1em;"><strong>타겟</strong>: 지자체 행정 서비스 제공자 (구청, 시청) 및 건물주</p>
        """, elem_classes="header-text")
        
        with gr.Tabs():
            # 탭 1: 패턴 만들기
            with gr.TabItem("🎨 보도블록 패턴 만들기"):
                gr.Markdown("### 원하는 보도블록 패턴을 텍스트로 설명해주세요")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        pattern_prompt = gr.Textbox(
                            label="패턴 설명 (색깔, 모양, 질감 등을 자세히 설명해주세요)",
                            placeholder="예: 붉은색, 십자가 모양, 굴곡이 있는, 현대적인",
                            lines=3
                        )
                        
                        reference_image = gr.Image(
                            label="참고 이미지 (선택사항) - 참고할 패턴이나 디자인 이미지를 업로드하세요",
                            type="filepath"
                        )
                        
                        generate_btn = gr.Button("패턴 생성하기", variant="primary", size="lg")
                
                with gr.Row():
                    pattern1 = gr.Image(label="패턴 1", show_download_button=True)
                    pattern2 = gr.Image(label="패턴 2", show_download_button=True)
                
                pattern_status = gr.Textbox(label="상태", interactive=False)
                
                # 예시 프롬프트
                gr.Examples(
                    examples=[
                        ["빨간색과 회색이 교차하는 체크무늬 패턴"],
                        ["육각형 모양의 현대적인 디자인, 그레이 톤"],
                        ["전통적인 벽돌 패턴, 따뜻한 갈색"],
                        ["물결 무늬가 있는 블루 톤의 미끄럼 방지 블록"]
                    ],
                    inputs=[pattern_prompt]
                )
            
            # 탭 2: 패턴 적용하기  
            with gr.TabItem("🛣️ 보도블록 패턴 적용하기"):
                gr.Markdown("### 도로 사진에 새로운 보도블록 패턴을 적용해보세요")
                
                with gr.Row():
                    with gr.Column():
                        road_image = gr.Image(
                            label="도로/인도 사진 - 보도블록이 보이는 도로나 인도 사진을 업로드하세요",
                            type="filepath"
                        )
                        
                        pattern_ref = gr.Image(
                            label="적용할 패턴 - 위에서 생성한 패턴이나 다른 패턴 이미지를 업로드하세요",
                            type="filepath"
                        )
                        
                        apply_btn = gr.Button("패턴 적용하기", variant="primary", size="lg")
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("#### Before (원본)")
                        original_display = gr.Image(label="원본 도로", interactive=False)
                    
                    with gr.Column():
                        gr.Markdown("#### After (패턴 적용)")
                        result_image = gr.Image(label="결과", show_download_button=True)
                
                apply_status = gr.Textbox(label="상태 (before와 after가 같은 경우 다시 실행해주세요)", interactive=False)
        
        # 이벤트 핸들러
        generate_btn.click(
            fn=generate_pattern,
            inputs=[pattern_prompt, reference_image],
            outputs=[pattern1, pattern2, pattern_status]
        )
        
        apply_btn.click(
            fn=apply_pattern_to_road,
            inputs=[road_image, pattern_ref],
            outputs=[result_image, apply_status]
        )
        
        # 원본 이미지 미리보기
        road_image.change(
            fn=lambda x: x,
            inputs=[road_image], 
            outputs=[original_display]
        )

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