app.py

import spaces
import torch
from diffusers import FluxPipeline
import gradio as gr
import random
import numpy as np
import os

if torch.cuda.is_available():
    device = "cuda"
    print("GPU를 사용합니다")
else:
    device = "cpu"
    print("CPU를 사용합니다")

HF_TOKEN = os.getenv("HF_TOKEN")
MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.to(device)

@spaces.GPU(duration=160)
def generate_image(prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt, progress=gr.Progress(track_tqdm=True)):
    if seed is None or seed == 0:
        seed = random.randint(1, MAX_SEED)
    
    generator = torch.Generator().manual_seed(seed)
    
    with torch.inference_mode():
        output = pipe(
            prompt=prompt,
            num_inference_steps=num_inference_steps,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            generator=generator,
            num_images_per_prompt=num_images_per_prompt
        ).images
    
    return output

def random_seed():
    return random.randint(1, MAX_SEED)

def create_random_seed():
    new_seed = random_seed()
    return [gr.Number.update(value=new_seed), f"현재 시드: {new_seed}"]

examples = [
    ["A cat holding a sign that says hello world"],
    ["a tiny astronaut hatching from an egg on the moon"],
    ["An astronaut on mars in a futuristic cyborg suit"],
]

css = """
.gradio-container {
    max-width: 1400px !important;
    margin: auto;
}
.image-container img {
    max-height: 600px !important;
}
.image-slider {
    height: 600px !important;
    max-height: 600px !important;
}
h1 {
    text-align: center;
    font-family: 'Pretendard', sans-serif;
    color: #EA580C;
    font-size: 2.5rem;
    font-weight: 700;
    margin-bottom: 1.5rem;
    text-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.subtitle {
    text-align: center;
    color: #4B5563;
    font-size: 1.1rem;
    margin-bottom: 2rem;
    font-family: 'Pretendard', sans-serif;
}
.gr-button-primary {
    background-color: #F97316 !important;
    border: none !important;
    box-shadow: 0 2px 4px rgba(234, 88, 12, 0.2) !important;
}
.gr-button-primary:hover {
    background-color: #EA580C !important;
    transform: translateY(-1px);
    box-shadow: 0 4px 6px rgba(234, 88, 12, 0.25) !important;
}
.footer-content {
    text-align: center;
    margin-top: 3rem;
    padding: 2rem;
    background: linear-gradient(to bottom, #FFF7ED, white);
    border-radius: 12px;
    font-family: 'Pretendard', sans-serif;
}
.footer-content a {
    color: #EA580C;
    text-decoration: none;
    font-weight: 500;
    transition: all 0.2s;
}
.footer-content a:hover {
    color: #C2410C;
}
.visit-button {
    background-color: #EA580C;
    color: white !important;
    padding: 12px 24px;
    border-radius: 8px;
    font-weight: 600;
    text-decoration: none;
    display: inline-block;
    transition: all 0.3s;
    margin-top: 1rem;
    box-shadow: 0 2px 4px rgba(234, 88, 12, 0.2);
    font-size: 1.1rem;
}
.visit-button:hover {
    background-color: #C2410C;
    transform: translateY(-2px);
    box-shadow: 0 4px 6px rgba(234, 88, 12, 0.25);
    color: white !important;
}
.container-wrapper {
    background: white;
    border-radius: 16px;
    padding: 2rem;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
}
.image-container {
    border-radius: 12px;
    overflow: hidden;
    border: 2px solid #F3F4F6;
}
.seed-button {
    background-color: #F3F4F6 !important;
    color: #374151 !important;
    border: 1px solid #E5E7EB !important;
    border-radius: 8px !important;
    padding: 8px 16px !important;
    font-size: 0.9rem !important;
    font-weight: 500 !important;
    transition: all 0.2s ease-in-out !important;
}
.seed-button:hover {
    background-color: #E5E7EB !important;
    transform: translateY(-1px);
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
}
.seed-text {
    font-family: 'Pretendard', sans-serif;
    color: #4B5563;
    font-size: 0.9rem;
    margin: 0;
    line-height: 2.5;
    font-weight: 500;
}
"""

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue=gr.themes.Color(
            c50="#FFF7ED",
            c100="#FFEDD5",
            c200="#FED7AA",
            c300="#FDBA74",
            c400="#FB923C",
            c500="#F97316",
            c600="#EA580C",
            c700="#C2410C",
            c800="#9A3412",
            c900="#7C2D12",
            c950="#431407",
        ),
        secondary_hue="zinc",
        neutral_hue="zinc",
        font=("Pretendard", "sans-serif")
    ),
    css=css
) as demo:
    with gr.Row():
        with gr.Column():
            gr.HTML(
                """
                <h1>끝장AI FLUX.1 이미지 생성기</h1>
                <div class="subtitle">
                    강력한 AI 기술로 당신의 상상을 현실로 만들어보세요
                </div>
                """
            )
    
    with gr.Group(elem_classes="container-wrapper"):
        with gr.Column():
            prompt = gr.Textbox(
                label="프롬프트",
                info="원하는 이미지를 설명해주세요",
                placeholder="고양이..."
            )
            run_button = gr.Button("생성하기", variant="primary")
            result = gr.Gallery(
                label="생성된 AI 이미지",
                elem_id="gallery",
                elem_classes="image-container"
            )
    
    with gr.Accordion("고급 설정", open=False):
        with gr.Row():
            num_inference_steps = gr.Slider(
                label="추론 단계 수",
                info="이미지의 디노이징 단계 수입니다. 더 많은 단계는 더 높은 품질의 이미지를 생성하지만 시간이 더 걸립니다",
                minimum=1,
                maximum=50,
                value=25,
                step=1,
                interactive=True,
                show_label=True,
                container=True,
                randomize=False
            )
            guidance_scale = gr.Slider(
                label="가이던스 스케일",
                info="텍스트 프롬프트를 얼마나 충실히 따를지 제어합니다. 높은 값은 입력 텍스트에 더 가깝게 생성됩니다",
                minimum=0.0,
                maximum=7.0,
                value=3.5,
                step=0.1,
                interactive=True,
                show_label=True,
                container=True,
                randomize=False
            )
        
        with gr.Row():
            width = gr.Slider(
                label="너비",
                info="이미지의 너비",
                minimum=256,
                maximum=1024,
                step=32,
                value=1024,
                interactive=True,
                show_label=True,
                container=True,
                randomize=False
            )
            height = gr.Slider(
                label="높이",
                info="이미지의 높이",
                minimum=256,
                maximum=1024,
                step=32,
                value=1024,
                interactive=True,
                show_label=True,
                container=True,
                randomize=False
            )
        
        with gr.Row():
            seed = gr.Number(value=random_seed(), minimum=0, maximum=MAX_SEED, step=1, visible=False)
            seed_button = gr.Button("🔄 새로운 시드 생성", size="sm")
            seed_text = gr.Markdown("")  # 현재 시드값 표시
            num_images_per_prompt = gr.Slider(
                label="프롬프트당 이미지 수",
                info="설정된 값으로 생성할 이미지의 수",
                minimum=1,
                maximum=4,
                step=1,
                value=2
            )
        
        seed_button.click(fn=create_random_seed, outputs=[seed, seed_text])

# Add after other gr.on() events:
    seed_button.click(
        fn=create_random_seed,
        outputs=[seed, seed_text]
    )

    gr.Examples(
        examples=examples,
        fn=generate_image,
        inputs=[prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt],
        outputs=[result],
        cache_examples=CACHE_EXAMPLES
    )

    seed_button.click(
        fn=create_random_seed,
        outputs=[seed, seed_text]
    )

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

    gr.HTML(
        """
        <div class="footer-content">
            <p style="font-size: 1.1rem; font-weight: 500; color: #1F2937;">끝장AI가 제공하는 고급 AI 도구를 더 경험하고 싶으신가요?</p>
            <a href="https://finalendai.com" target="_blank" class="visit-button">
                끝장AI 방문하기
            </a>
            <p style="margin-top: 1.5rem; color: #6B7280; font-size: 0.9rem;">
                © 2024 끝장AI. All rights reserved.
            </p>
        </div>
        """
    )

demo.queue().launch(share=False)