import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline
import spaces
# 기본 설정
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

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

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

# 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
]

@spaces.GPU()  # 이 데코레이터를 다시 추가
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=0.0
    ).images[0]
    return image, seed

# CSS 스타일 수정
css = """
.container {
    display: flex;
    flex-direction: row;
    height: 100%;
}

.input-column {
    flex: 1;
    padding: 20px;
    border-right: 2px solid #eee;
    max-width: 800px;  /* 증가된 최대 너비 */
}

.examples-column {
    flex: 1;
    padding: 20px;
    overflow-y: auto;
    background: #f7f7f7;
}

.title {
    text-align: center;
    color: #2a2a2a;
    padding: 20px;
    font-size: 2.5em;
    font-weight: bold;
    background: linear-gradient(90deg, #f0f0f0 0%, #ffffff 100%);
    border-bottom: 3px solid #ddd;
    margin-bottom: 30px;
}

.subtitle {
    text-align: center;
    color: #666;
    margin-bottom: 30px;
}

.input-box {
    background: white;
    padding: 20px;
    border-radius: 10px;
    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
    margin-bottom: 20px;
    width: 100%;  /* 전체 너비 사용 */
}

/* 텍스트 입력 박스 스타일 */
.input-box textarea {
    width: 100% !important;  /* 강제 너비 100% */
    min-width: 600px !important;  /* 최소 너비 설정 */
    font-size: 14px !important;
    line-height: 1.5 !important;
    padding: 12px !important;
}

.example-card {
    background: white;
    padding: 15px;
    margin: 10px 0;
    border-radius: 8px;
    box-shadow: 0 2px 5px rgba(0,0,0,0.05);
}

.example-title {
    font-weight: bold;
    color: #2a2a2a;
    margin-bottom: 10px;
}

/* 레이아웃 조정 */
.contain {
    max-width: 1400px !important;  /* 전체 컨테이너 너비 증가 */
    margin: 0 auto !important;
}

/* 입력 영역 조정 */
.input-area {
    flex: 2 !important;  /* 입력 영역 비율 증가 */
}

/* 예제 영역 조정 */
.examples-area {
    flex: 1 !important;
}
"""


# Gradio 인터페이스 수정
with gr.Blocks(css=css) as demo:
    gr.Markdown(
        """
        <div class="title">FLUX Mindmap Generator</div>
        <div class="subtitle">Create beautiful hand-drawn style diagrams using FLUX AI</div>
        """)


    
    with gr.Row(equal_height=True) as main_row:
        # 왼쪽 입력 컬럼
        with gr.Column(elem_id="input-column", scale=2):  # scale 값 증가
            with gr.Group(elem_classes="input-box"):
                prompt = gr.Text(
                    label="Diagram Prompt",
                    placeholder="Enter your diagram structure...",
                    lines=10,  # 줄 수 증가
                    elem_classes="prompt-input"  # CSS 클래스 추가
                )
   
                run_button = gr.Button("Generate Diagram", variant="primary")
                result = gr.Image(label="Generated Diagram")
            
            with gr.Accordion("Advanced Settings", open=False):
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,
                    )
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,
                    )
                
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        # 오른쪽 예제 컬럼
        with gr.Column(elem_id="examples-column", scale=1):
            gr.Markdown("### Example Diagrams")
            for example in EXAMPLES:
                with gr.Group(elem_classes="example-card"):
                    gr.Markdown(f"#### {example['title']}")
                    gr.Markdown(f"```\n{example['prompt']}\n```")
                    
                    def create_example_handler(ex):
                        def handler():
                            return {
                                prompt: ex["prompt"],
                                width: ex["width"],
                                height: ex["height"]
                            }
                        return handler
                    
                    gr.Button("Use This Example", size="sm").click(
                        fn=create_example_handler(example),
                        outputs=[prompt, width, height]
                    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
        outputs=[result, seed]
    )

# 앱 실행
if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )