import gradio as gr
from PIL import Image
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration
from diffusers import StableDiffusionPipeline
import traceback

# 모델 로딩 (최대한 빠른 처리: BLIP + SD-Turbo)
device = "cuda" if torch.cuda.is_available() else "cpu"

blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)

pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-turbo",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to(device)
pipe.safety_checker = None  # 안전필터 off (실습/학습 목적)

def caricature_from_photo(image):
    try:
        raw_image = image.convert('RGB')
        inputs = blip_processor(raw_image, return_tensors="pt").to(device)
        out = blip_model.generate(**inputs, max_new_tokens=30)
        description = blip_processor.decode(out[0], skip_special_tokens=True)

        # 캐리커쳐 스타일 프롬프트 (이모티콘과 직관적 키워드로 꾸밈)
        caricature_prompt = (
            f"{description}, cute cartoon, gentle lines, charming and lovely, sweet expression, lovely mood, digital illustration, "
            "cheerful, simple background, joy, cartoon-style, 😄✨"
        )

        # turbo SD: 매우 빠른 생성!
        result = pipe(
            caricature_prompt, 
            num_inference_steps=4,         # turbo 권장값: 4
            guidance_scale=0.0              # turbo 권장값: 0.0 (텍스트 프롬프트 반영)
        ).images[0]

        info_text = (
            f"**이미지 설명**: {description}\n\n"
            f"**캐리커쳐 프롬프트**: `{caricature_prompt}`"
        )
        return result, info_text
    except Exception as e:
        err = f"❗️ 오류가 발생했어요!\n\n{e}\n\n{traceback.format_exc()}"
        empty = Image.new("RGB", (512,512), "#FFFDE7")
        return empty, err

custom_css = """
body { background: linear-gradient(135deg, #a7ffeb 0%, #fce4ec 100%);}
.gradio-container { background-color: #fafcff !important;}
.gr-button { background: #ffd600; color: #252525; border-radius:18px; font-size:1.2em;}
.gr-button:hover { background:#ffd54f;}
#caricature-preview img {
    border-radius: 30px; border: 4px solid #ffff8d;
    box-shadow: 0px 8px 56px #ffd60044;
    margin-top: 1em;
}
#main-view .block {
    background: #ffffffb8; border-radius: 22px;
    box-shadow: 0 4px 32px #b3afec33; margin:18px;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    <div style="text-align:center; padding-top:.5em;">
      <span style="font-size:2.3em;">🎨 <b style="color:#1ec69d">캐리커쳐 매직</b> <span style="font-size:1em;">앱</span> ✨</span>
      <p style="font-size:1.11em; color:#575d76;">
        어떤 사진이든 <b>1초 만에 캐리커쳐</b>로!<br>
        <span style="color:#2196f3">카메라 📷</span> 또는 <span style="color:#ff1696">앨범</span>에서 이미지를 올려보세요!
      </p>
    </div>
    """)
    with gr.Row(elem_id="main-view"):
        with gr.Column():
            photo = gr.Image(
                label="📷 사진 찍기 또는 선택(앨범)",
                type="pil",
                height=300,
                elem_id="user-photo",
            )
            btn = gr.Button("👉 캐리커쳐 만들기!", elem_id="do_caricature")
        with gr.Column():
            caricature_img = gr.Image(
                label="🖼️ 캐리커쳐 미리보기",
                elem_id="caricature-preview",
                height=350,
                show_label=True
            )
            info = gr.Markdown(label="생성 내역 및 설명", value="", elem_id="summary-field")
    btn.click(caricature_from_photo, inputs=photo, outputs=[caricature_img, info])

    # 안내 메시지
    gr.Markdown("""
    <div style="background:#fffde7; border-radius:18px; border:1.5px dashed #ffe082; padding:1em; text-align:center; margin-top:18px;">
      <b>💡 사용법 & 인쇄 안내</b><br/>
      <ul style="text-align:left; margin-left:1.8em;">
        <li>휴대폰은 [사진찍기] 선택 시, 카메라와 앨범 둘 다 사용 가능</li>
        <li>[👉 캐리커쳐 만들기!] 버튼을 누르면 캐리커쳐가 빠르게 생성됩니다.</li>
        <li>결과 이미지를 길게 누르거나 우클릭하여 저장/공유/인쇄 할 수 있습니다.</li>
      </ul>
    </div>
    """)

demo.launch()