app.py

import os
import random
import numpy as np
import gradio as gr

# --- Hugging Face Spaces 데코레이터(없어도 동작하도록 대체) ---
try:
    import spaces
    GPU_DECORATOR = spaces.GPU
except Exception:
    def GPU_DECORATOR(*args, **kwargs):
        def _wrap(fn):
            return fn
        return _wrap

import torch
from diffusers import DiffusionPipeline

# -----------------------
# 장치/정밀도 자동 선택
# -----------------------
def select_device_dtype():
    if torch.cuda.is_available():
        device = "cuda"
        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    else:
        device = "cpu"
        dtype = torch.float32
    return device, dtype

device, dtype = select_device_dtype()
if device == "cuda":
    torch.backends.cuda.matmul.allow_tf32 = True

# 제한값
MAX_SEED = np.iinfo(np.uint32).max
MAX_SIDE = 2048
MAX_TOTAL_PIXELS = 2_359_296  # 해상도 총 픽셀 가드(약 2.36MP)

# -----------------------
# 모델 로드 (FLUX.1-schnell)
# -----------------------
MODEL_ID = "black-forest-labs/FLUX.1-schnell"
pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype).to(device)
if hasattr(pipe, "enable_vae_slicing"):
    pipe.enable_vae_slicing()
if hasattr(pipe, "enable_vae_tiling"):
    pipe.enable_vae_tiling()

# -----------------------
# 유틸
# -----------------------
def _clamp_hw(width: int, height: int):
    width = int(max(256, min(int(width), MAX_SIDE)))
    height = int(max(256, min(int(height), MAX_SIDE)))
    if width * height > MAX_TOTAL_PIXELS:
        scale = (MAX_TOTAL_PIXELS / (width * height)) ** 0.5
        width = int((width * scale) // 32 * 32)
        height = int((height * scale) // 32 * 32)
        width = max(256, min(width, MAX_SIDE))
        height = max(256, min(height, MAX_SIDE))
    return width, height

def _validate_params(prompt: str, steps: int, guidance: float):
    if not prompt or not prompt.strip():
        raise ValueError("프롬프트가 비어 있습니다. 내용을 입력해 주세요.")
    if not (1 <= steps <= 50):
        raise ValueError("추론 스텝은 1~50 범위여야 합니다.")
    if not (0.0 <= guidance <= 20.0):
        raise ValueError("가이던스 스케일은 0.0~20.0 범위여야 합니다.")

# -----------------------
# 생성 함수
# -----------------------
@GPU_DECORATOR()
def generate_image(prompt, seed, randomize_seed, width, height, steps, guidance_scale, progress=gr.Progress(track_tqdm=True)):
    try:
        prompt = prompt.strip()
        _validate_params(prompt, steps, guidance_scale)

        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
        generator = torch.Generator(device=device).manual_seed(int(seed))

        width, height = _clamp_hw(width, height)

        progress(0.1, desc="초기화 중…")
        if device == "cuda":
            autocast_ctx = torch.autocast(device_type="cuda", dtype=dtype)
        elif device == "cpu":
            autocast_ctx = torch.autocast(device_type="cpu", dtype=dtype) if dtype != torch.float32 else torch.no_grad()
        else:
            autocast_ctx = torch.no_grad()

        with autocast_ctx:
            progress(0.4, desc="이미지 생성 중…")
            out = pipe(
                prompt=prompt,
                width=int(width),
                height=int(height),
                num_inference_steps=int(steps),
                generator=generator,
                guidance_scale=float(guidance_scale)
            )
            image = out.images[0]

        progress(1.0, desc="완료")
        return image, int(seed)
    except Exception as e:
        gr.Error(f"생성 중 오류: {type(e).__name__}: {e}")
        return None, int(seed)

def set_prompt(example_text):
    return example_text

# -----------------------
# 한글 예시 프롬프트
# -----------------------
example_prompts = {
    "제품 디자인": [
        """슬릭한 인더스트리얼 스타일의 커피머신 컨셉 스케치:
- 곡선형 메탈 바디, 베젤 최소화
- 설정용 터치스크린 패널
- 모던한 매트 블랙 마감
- 손그림(스케치) 컨셉 아트 스타일"""
    ],
    "마인드맵": [
        """손그림 스타일의 다채로운 마인드맵, 교육용, 생동감 있는 색감, 명확한 계층 구조, 황금비 레이아웃.
KNOWLEDGE
├── ACQUISITION [번개가 친 뇌 ~60px]
│   ├── READING [빛나는 펼친 책]
│   ├── PRACTICE [공구 아이콘]
│   └── OBSERVATION [확대경 든 눈]
├── PROCESSING [기어 네트워크 ~50px]
│   ├── ANALYSIS [상승 그래프]
│   └── SYNTHESIS [퍼즐 조각]
├── RETENTION [메모리 칩 ~45px]
│   ├── SHORT-TERM [번쩍임]
│   └── LONG-TERM [견고한 아카이브]
└── APPLICATION
    ├── CREATION [팔레트]
    └── INNOVATION [별자리 전구]"""
    ],
    "목업": [
        """손그림 와이어프레임 스타일의 모바일 뱅킹 앱 목업:
- 로고가 있는 타이틀 화면
- 로그인(아이디, 비밀번호, 로그인 버튼)
- 대시보드 3개 섹션(잔액, 거래내역, 빠른 실행)
- 하단 내비게이션(홈, 이체, 프로필)"""
    ],
    "인포그래픽": [
        """대기업 연차보고서용 플랫 스타일 인포그래픽:
- 제목: "Global Renewable Energy Trends 2025"
- 부제: "시장점유율과 성장분석"
- 시각 요소:
    - 지역별 태양광·풍력·수력 생산량 막대그래프
    - 에너지 비중 파이차트: 태양광(45%)·풍력(30%)·수력(25%)
    - 연도별 성장 추세선
- 아이콘: 미니멀 태양, 풍력 터빈, 물방울
- 레이아웃: 그리드 기반, 파스텔 포인트, 화이트 스페이스 충분
- 주석: KPI 핵심 수치 및 전망 콜아웃"""
    ],
    "다이어그램": [
        """엔드투엔드 비즈니스 워크플로 다이어그램(손그림, 교육적·전문적):
- 제목: "통합 비즈니스 프로세스"
- 구성:
    - 시장분석(차트, 경쟁자 맵)
    - 전략수립(브레인스토밍 클라우드, 핵심 포커스)
    - 제품설계(스케치, 피드백 루프)
    - 구현(타임라인 마커, 리소스 아이콘)
    - 출시 후 리뷰(지표, 지속 개선)
- 명확한 방향 화살표, 색상으로 단계 구분"""
    ],
    "플로우차트": [
        """손그림 스타일 플로우차트, 선명한 색, 미니멀 아이콘.
BUSINESS WORKFLOW
├── 시작 [초록 버튼 ~40px]
│   ├── 요구사항 수집 [폴더 아이콘]
│   └── 데이터 분석 [차트 아이콘]
├── 구현 [코딩 심볼 ~50px]
│   ├── 프론트엔드 [브라우저 아이콘]
│   └── 백엔드 [서버 아이콘]
├── 테스트 & 통합 [기어 아이콘 ~45px]
└── 배포
    └── 종료 [체커기 깃발 ~40px]"""
    ]
}

# -----------------------
# Gradio UI (한국어)
# -----------------------
css = """
* { box-sizing: border-box; }
body {
    background: linear-gradient(135deg, #667eea, #764ba2);
    font-family: 'Pretendard', 'Apple SD Gothic Neo', 'Noto Sans KR', 'Helvetica Neue', Arial, sans-serif;
    color: #333; margin: 0; padding: 0;
}
.gradio-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 30px 40px;
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
    margin: 40px auto;
    width: 1200px;
    overflow: visible !important;
}
.sidebar {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
    position: relative; z-index: 10;
    overflow: visible !important;
}
button, .btn {
    background: linear-gradient(90deg, #ff8a00, #e52e71);
    border: none; color: #fff;
    padding: 12px 24px; text-transform: uppercase;
    font-weight: bold; letter-spacing: 1px;
    border-radius: 5px; cursor: pointer;
    transition: transform 0.2s ease-in-out;
}
button:hover, .btn:hover { transform: scale(1.05); }
.example-accordion { width: 100% !important; max-width: 100% !important; }
.example-accordion button { width: auto !important; white-space: normal !important; }
"""

with gr.Blocks(css=css, title="워크플로 캔버스") as demo:
    gr.Markdown(
        """
        <div style="text-align:center;">
            <h1>워크플로 캔버스</h1>
            <p>여러 탭에서 비즈니스에 필요한 디자인 컨셉과 워크플로 다이어그램을 생성해 보세요.</p>
            <p><strong>커뮤니티:</strong> <a href="https://discord.gg/openfreeai" target="_blank">https://discord.gg/openfreeai</a></p>
        </div>
        """
    )
    gr.HTML(
        """<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fginigen-Workflow-Canvas.hf.space">
        <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fginigen-Workflow-Canvas.hf.space&countColor=%23263759" alt="방문자 배지"/>
        </a>"""
    )

    with gr.Row():
        # 왼쪽 사이드바: 공통 파라미터
        with gr.Column(scale=2, elem_classes="sidebar"):
            gr.Markdown("### 생성 파라미터")

            size_preset = gr.Dropdown(
                label="해상도 프리셋",
                choices=[
                    "1024x1024", "1536x1024", "1024x1536",
                    "1344x1344", "1536x1536", "1920x1080", "1080x1920"
                ],
                value="1024x1024"
            )
            width_slider = gr.Slider(label="가로 해상도(px)", minimum=256, maximum=MAX_SIDE, step=32, value=1024)
            height_slider = gr.Slider(label="세로 해상도(px)", minimum=256, maximum=MAX_SIDE, step=32, value=1024)

            def apply_preset(preset):
                w, h = map(int, preset.split("x"))
                w, h = _clamp_hw(w, h)
                return w, h
            size_preset.change(fn=apply_preset, inputs=size_preset, outputs=[width_slider, height_slider])

            seed_slider = gr.Slider(label="시드(Seed)", minimum=0, maximum=int(MAX_SEED), step=1, value=42)
            randomize_seed = gr.Checkbox(label="시드 랜덤", value=True)

            def toggle_seed(disable):
                return gr.update(interactive=not disable)
            randomize_seed.change(fn=toggle_seed, inputs=randomize_seed, outputs=seed_slider)

            steps_slider = gr.Slider(label="추론 스텝", minimum=1, maximum=50, step=1, value=20)
            guidance_slider = gr.Slider(label="가이던스 스케일", minimum=0.0, maximum=20.0, step=0.5, value=7.5)

        # 메인: 탭 UI
        with gr.Column(scale=8):
            with gr.Tabs():
                def build_tab(tab_title, ex_key, placeholder_text):
                    with gr.Tab(tab_title):
                        tb = gr.Textbox(
                            label=f"{tab_title} 프롬프트",
                            placeholder=placeholder_text,
                            lines=5,
                            value=example_prompts[ex_key][0]
                        )
                        btn = gr.Button(f"{tab_title} 생성")
                        img = gr.Image(label=f"{tab_title} 결과", type="pil", value=None, height=512)
                        with gr.Accordion("예시 프롬프트", open=True, elem_classes="example-accordion"):
                            for ex in example_prompts[ex_key]:
                                gr.Button(ex, variant="secondary").click(
                                    fn=lambda ex=ex: set_prompt(ex),
                                    outputs=tb
                                )
                        btn.click(
                            fn=generate_image,
                            inputs=[tb, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
                            outputs=[img, seed_slider]
                        )

                build_tab("제품 디자인", "제품 디자인", "제품 디자인 컨셉을 입력하세요…")
                build_tab("마인드맵", "마인드맵", "마인드맵 설명을 입력하세요…")
                build_tab("목업", "목업", "앱/웹 목업 설명을 입력하세요…")
                build_tab("인포그래픽", "인포그래픽", "인포그래픽 설명을 입력하세요…")
                build_tab("다이어그램", "다이어그램", "다이어그램 설명을 입력하세요…")
                build_tab("플로우차트", "플로우차트", "플로우차트 설명을 입력하세요…")

if __name__ == "__main__":
    demo.queue(concurrency_count=2, max_size=32)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )