app.py

import gradio as gr
from gradio_client import Client
import os
from dotenv import load_dotenv
import warnings

# 경고 메시지 숨기기
warnings.filterwarnings('ignore', category=UserWarning)

# 환경 변수 로드
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

def generate_diagram(prompt, width=1024, height=1024):
    """FLUX AI를 사용하여 다이어그램 생성"""
    try:
        # API 호출을 통해 이미지 생성
        result = gr.Image.update(value=None)  # 초기 상태는 빈 이미지
        
        # Hugging Face API 호출
        client = Client(
            "https://black-forest-labs-flux-1-schnell.hf.space",  # API 엔드포인트 직접 지정
            hf_token=HF_TOKEN,
        )
        
        # 이미지 생성
        result = client.predict(
            prompt,
            1872187377,  # seed
            False,       # randomize_seed
            width,
            height,
            4,          # num_inference_steps
            api_name="/infer"
        )
        return result
    except Exception as e:
        raise gr.Error(f"다이어그램 생성 중 오류 발생: {str(e)}")

# Convert example format for Gradio
GRADIO_EXAMPLES = [
    [
        """A handrawn colorful mind map diagram, educational style, vibrant colors, clear hierarchy.
        KNOWLEDGE
        ├── ACQUISITION [Brain with Lightning ~60px]
        │   ├── READING [Open Book with Glow]
        │   └── PRACTICE [Hands-on Tools]
        └── APPLICATION
            ├── CREATION [Artist Palette]
            └── INNOVATION [Lightbulb]""",
        1024,
        1024
    ],
    [
        """A handrawn colorful mind map diagram, tech-focused style, neon accents.
        DIGITAL TRANSFORM
        ├── CLOUD [Cloud with Data ~55px]
        │   ├── STORAGE [Database]
        │   └── COMPUTING [Server]
        └── SECURITY
            ├── ENCRYPTION [Lock]
            └── MONITORING [Radar]""",
        1024,
        1024
    ],
    [
        """A handrawn colorful mind map diagram, creative style, flowing design.
        INNOVATION
        ├── IDEATION [Lightbulb ~60px]
        │   ├── RESEARCH [Magnifier]
        │   └── BRAINSTORM [Brain]
        └── EXECUTION
            ├── PROTOTYPE [Tools]
            └── TEST [Checklist]""",
        1024,
        1024
    ]
]


# Enhanced examples with more detailed prompts and specific styling
EXAMPLES = [
    {
        "title": "Knowledge Tree",
        "prompt": """A handrawn colorful mind map diagram, educational style, vibrant colors, clear hierarchy, golden ratio layout.
        KNOWLEDGE
        ├── ACQUISITION [Brain with Lightning ~60px]
        │   ├── READING [Open Book with Glow]
        │   ├── PRACTICE [Hands-on Tools]
        │   └── OBSERVATION [Eye with Magnifier]
        ├── PROCESSING [Gear Network ~50px]
        │   ├── ANALYSIS [Graph Trending Up]
        │   └── SYNTHESIS [Puzzle Pieces]
        ├── RETENTION [Memory Chip ~45px]
        │   ├── SHORT-TERM [Quick Flash]
        │   └── LONG-TERM [Solid Archive]
        └── APPLICATION
            ├── CREATION [Artist Palette]
            └── INNOVATION [Lightbulb Constellation]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Digital Transformation",
        "prompt": """A handrawn colorful mind map diagram, tech-focused style, neon accents, circuit board patterns.
        DIGITAL TRANSFORM
        ├── CLOUD [Cloud with Data Rain ~55px]
        │   ├── STORAGE [Database Cluster]
        │   └── COMPUTING [Server Array]
        ├── AUTOMATION [Robot Arm ~50px]
        │   ├── WORKFLOWS [Flowchart]
        │   └── AI/ML [Neural Network]
        ├── SECURITY [Shield Matrix ~45px]
        │   ├── ENCRYPTION [Lock Code]
        │   └── MONITORING [Radar Screen]
        └── INTEGRATION
            ├── APIS [Puzzle Connect]
            └── MICROSERVICES [Building Blocks]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Creative Process",
        "prompt": """A handrawn colorful mind map diagram, artistic style, watercolor effects, flowing connections.
        CREATIVITY
        ├── INSPIRATION [Constellation Stars ~60px]
        │   ├── NATURE [Organic Patterns]
        │   └── CULTURE [Global Icons]
        ├── IDEATION [Floating Bubbles ~50px]
        │   ├── BRAINSTORM [Thunder Cloud]
        │   └── REFINEMENT [Diamond Polish]
        ├── EXECUTION [Artist Tools ~45px]
        │   ├── TECHNIQUE [Skilled Hands]
        │   └── MEDIUM [Palette Mix]
        └── PRESENTATION
            ├── GALLERY [Frame Display]
            └── FEEDBACK [Echo Ripples]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Future Cities",
        "prompt": """A handrawn colorful mind map diagram, futuristic style, holographic elements, sustainable themes.
        SMART CITY
        ├── MOBILITY [Hover Transport ~60px]
        │   ├── AUTONOMOUS [Self-Driving]
        │   └── CONNECTED [Network Grid]
        ├── ENERGY [Solar Crystal ~55px]
        │   ├── RENEWABLE [Green Power]
        │   └── STORAGE [Battery Hub]
        ├── LIVING [Eco Building ~50px]
        │   ├── VERTICAL [Sky Gardens]
        │   └── COMMUNITY [People Connect]
        └── INFRASTRUCTURE
            ├── AI GRID [Neural City]
            └── ECO SYSTEM [Nature Tech]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Health Evolution",
        "prompt": """A handrawn colorful mind map diagram, medical style, DNA helix patterns, wellness focus.
        HEALTH 3.0
        ├── PREVENTION [Shield DNA ~60px]
        │   ├── LIFESTYLE [Activity Pulse]
        │   └── MONITORING [Health Watch]
        ├── TREATMENT [Caduceus Tech ~55px]
        │   ├── PERSONALIZED [DNA Code]
        │   └── REGENERATIVE [Cell Renew]
        ├── ENHANCEMENT [Upgrade Spiral ~50px]
        │   ├── COGNITIVE [Brain Boost]
        │   └── PHYSICAL [Body Optimize]
        └── INTEGRATION
            ├── AI HEALTH [Smart Doctor]
            └── COMMUNITY [Global Care]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Space Exploration",
        "prompt": """A handrawn colorful mind map diagram, cosmic style, star field background, planetary elements.
        SPACE FRONTIER
        ├── DISCOVERY [Telescope Array ~60px]
        │   ├── MAPPING [Star Charts]
        │   └── ANALYSIS [Data Stream]
        ├── TRAVEL [Rocket Launch ~55px]
        │   ├── PROPULSION [Energy Core]
        │   └── NAVIGATION [Space Map]
        ├── COLONIZATION [Dome City ~50px]
        │   ├── HABITATS [Life Sphere]
        │   └── RESOURCES [Mine Extract]
        └── RESEARCH
            ├── ASTROBIOLOGY [Life Search]
            └── PHYSICS [Space Time]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Ocean Innovation",
        "prompt": """A handrawn colorful mind map diagram, marine style, wave patterns, aqua themes.
        OCEAN TECH
        ├── EXPLORATION [Deep Submersible ~60px]
        │   ├── MAPPING [Sonar Wave]
        │   └── RESEARCH [Lab Bubble]
        ├── CONSERVATION [Marine Life ~55px]
        │   ├── PROTECTION [Reef Shield]
        │   └── RESTORATION [Growth Core]
        ├── HARVESTING [Sustainable Net ~50px]
        │   ├── ENERGY [Wave Power]
        │   └── RESOURCES [Bio Extract]
        └── MONITORING
            ├── AI SYSTEMS [Smart Sensors]
            └── ECOLOGY [Life Web]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Quantum Computing",
        "prompt": """A handrawn colorful mind map diagram, quantum style, wave-particle duality, matrix patterns.
        QUANTUM TECH
        ├── COMPUTATION [Qubit Matrix ~60px]
        │   ├── PROCESSING [Wave Function]
        │   └── ALGORITHMS [Code Quantum]
        ├── APPLICATIONS [Use Cases ~55px]
        │   ├── SIMULATION [Model World]
        │   └── OPTIMIZATION [Peak Find]
        ├── INFRASTRUCTURE [Q-Hardware ~50px]
        │   ├── CONTROL [Pulse Shape]
        │   └── COOLING [Zero Point]
        └── DEVELOPMENT
            ├── SOFTWARE [Q-Code Web]
            └── INTEGRATION [Classical Bridge]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "Bio Engineering",
        "prompt": """A handrawn colorful mind map diagram, biological style, DNA patterns, organic flow.
        BIOTECH
        ├── GENETICS [DNA Helix ~60px]
        │   ├── EDITING [CRISPR Tool]
        │   └── SYNTHESIS [Gene Build]
        ├── APPLICATIONS [Lab Array ~55px]
        │   ├── MEDICINE [Heal Cell]
        │   └── AGRICULTURE [Grow Plus]
        ├── PLATFORMS [Bio Factory ~50px]
        │   ├── SENSORS [Live Detect]
        │   └── PROCESSORS [Cell Compute]
        └── INTEGRATION
            ├── AI BIOLOGY [Smart Life]
            └── ECOSYSTEM [Nature Net]""",
        "width": 1024,
        "height": 1024
    },
    {
        "title": "AI Evolution",
        "prompt": """A handrawn colorful mind map diagram, neural network style, digital patterns, intelligence flow.
        AI FUTURE
        ├── COGNITION [Brain Network ~60px]
        │   ├── LEARNING [Growth Path]
        │   └── REASONING [Logic Tree]
        ├── PERCEPTION [Sensor Array ~55px]
        │   ├── VISION [Eye Matrix]
        │   └── LANGUAGE [Word Web]
        ├── INTERACTION [Connect Hub ~50px]
        │   ├── HUMAN [Bridge Link]
        │   └── MACHINE [Code Path]
        └── EVOLUTION
            ├── CONSCIOUSNESS [Mind Spark]
            └── CREATIVITY [Art Core]""",
        "width": 1024,
        "height": 1024
    }
]



# Convert examples to Gradio format
GRADIO_EXAMPLES = [
    [example["prompt"], example["width"], example["height"]]
    for example in EXAMPLES
]


demo = gr.Interface(
    fn=generate_diagram,
    inputs=[
        gr.Textbox(
            label="다이어그램 프롬프트",
            placeholder="다이어그램 구조를 입력하세요...",
            lines=10
        ),
        gr.Slider(
            label="너비",
            minimum=512,
            maximum=2048,
            step=128,
            value=1024
        ),
        gr.Slider(
            label="높이",
            minimum=512,
            maximum=2048,
            step=128,
            value=1024
        )
    ],
    outputs=gr.Image(label="생성된 다이어그램"),
    title="🎨 FLUX 다이어그램 생성기",
    description="FLUX AI를 사용하여 아름다운 손그림 스타일의 다이어그램을 생성합니다",
    article="""
    ### 더 나은 결과를 위한 팁
    - 명확한 계층 구조 사용
    - 대괄호 안에 아이콘 설명 포함
    - 간결하고 의미 있는 텍스트 사용
    - 일관된 형식 유지
    """,
    examples=GRADIO_EXAMPLES,
    cache_examples=True
)

# 앱 실행 부분만 수정
if __name__ == "__main__":
    demo.queue()  # 큐 기능을 별도로 활성화
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )