app.py

import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
from diffusers import DiffusionPipeline, AutoencoderTiny
from custom_pipeline import FluxWithCFGPipeline
from transformers import pipeline

# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 768
DEFAULT_INFERENCE_STEPS = 4

# Initialize translator
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# Device and model setup
dtype = torch.float16
pipe = FluxWithCFGPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
pipe.to("cuda")
torch.cuda.empty_cache()

# Translation function
def translate_to_english(text):
    if any(ord('가') <= ord(char) <= ord('힣') for char in text):
        translated = translator(text)[0]['translation_text']
        return translated
    return text

# Inference function
@spaces.GPU(duration=25)
def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, progress=gr.Progress(track_tqdm=True)):
    # Translate prompt if Korean
    english_prompt = translate_to_english(prompt)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(int(float(seed)))

    img = pipe.generate_images( 
            prompt=english_prompt,
            width=width,
            height=height,
            num_inference_steps=DEFAULT_INFERENCE_STEPS,
            generator=generator
        )
    return img, seed

# Example prompts
examples = [
"해변에서 밀짚모자를 쓰고 플로럴 드레스를 입은 여성이 'I Love'라고 적힌 작은 카드를 모래에 두고 있습니다. 황금빛 석양이 여유로운 해변 분위기를 자아냅니다.",    
"화려한 금발의 여성이 고급 호텔 스위트룸에서 포즈를 취하고 있으며, 스파클이 빛나는 황금빛 드레스를 입고 'I Love'라고 적힌 흰색 카드를 들고 있습니다. 부드러운 조명이 고급스러운 분위기를 자아냅니다.",
"근육질의 남성이 최신식 체육관에서 셔츠를 벗고 포즈를 취하고 있으며, 한 손에 단백질 쉐이크를 들고 있고 다른 손에는 'I Love '라고 적힌 카드가 있습니다. 밝고 깨끗한 조명이 그의 피지컬을 강조합니다.",
"자연스러운 해변 웨이브와 태닝된 피부를 가진 여성 여행 블로거가 열대 해변에서 황혼을 바라보며 앉아 있습니다. 흐르는 흰 드레스를 입고 있으며, 'I Love '라고 적힌 오래된 엽서를 들고 있습니다. 황금빛 조명이 따뜻한 분위기를 연출합니다.",
"도시 거리에서 완벽하게 스타일링된 머리와 수염을 가진 트렌디한 남성 패션 인플루언서가 세련된 정장을 입고 있으며, 'I Love '라고 적힌 검정색 명함을 들고 있습니다. 도심의 불빛이 세련된 분위기를 더합니다.",
"화장대에 앉아 자연스러운 메이크업을 한 젊은 뷰티 크리에이터가 파스텔 핑크색 가운을 입고 있으며, 'I Love '라고 적힌 메이크업 팔레트를 들고 있습니다. 부드러운 조명이 그녀의 피부를 빛나게 합니다.",
"네온사인으로 가득한 도시 야경 앞에 서 있는 스타일리시한 여성이 검은 가죽 재킷과 반짝이는 크롭탑을 입고 'I Love '라고 적힌 홀로그램 명함을 들고 있습니다. 사이버펑크 분위기의 화려한 네온 조명이 더해져 있습니다.",
"붉은 장미들 속에서 우아한 버건디 드레스를 입은 여성이 'I Love 라고 적힌 은색 카드를 들고 있습니다. 부드러운 조명이 그녀의 기품 있는 매력을 돋보이게 합니다.",
"고급 레스토랑에서 세련된 수트를 입은 남성이 와인 잔을 들고 있으며, 'I Love '라고 적힌 흰색 명함을 손에 들고 있습니다. 은은한 조명이 로맨틱한 분위기를 만듭니다.",
"빈티지 스타일의 서재 앞에 서 있는 여성이 체크무늬 코트를 입고 있으며, 'I Love '라고 적힌 옛 책을 들고 있습니다. 따뜻한 색감의 조명이 서재의 아늑한 느낌을 더합니다.",
"비 오는 밤의 어두운 골목길에서 회색 트렌치코트를 입은 남성이 'I Love '라고 적힌 검정색 카드를 들고 서 있습니다. 빗방울이 반사되어 긴장감 있는 분위기를 만듭니다.",
"고급 오피스 건물의 창가에 서 있는 남성이 슬림한 검정색 카드에 'I Love '라고 적힌 명함을 들고 있으며, 도시의 야경이 세련된 비즈니스 분위기를 연출합니다.",
"햇살이 비치는 화사한 정원에서 레이스 드레스를 입은 여성이 'I Love '라고 적힌 작은 손수건을 들고 있습니다. 부드러운 햇살이 따뜻한 분위기를 만들어줍니다.",
"푸른 호수를 배경으로 캐주얼하게 서 있는 남성이 청바지와 셔츠를 입고 'I Love '라고 적힌 나무판자를 들고 있습니다. 자연의 맑은 날씨가 상쾌함을 더해줍니다.",
"다채로운 조명 아래 화려한 디저트 테이블 앞에 서 있는 여성이 로맨틱한 드레스를 입고 있으며, 'I Love'라고 적힌 작은 케이크를 들고 있습니다. 축제 분위기를 더합니다.",
"바람이 휘날리는 바다 절벽 위에서 자유롭게 서 있는 남성이 흰 셔츠를 입고 있으며, 'I Love '라고 적힌 조약돌을 들고 있습니다. 하늘과 바다가 상쾌한 분위기를 연출합니다.",
"예술 갤러리에서 붉은 립스틱을 바르고 화려한 드레스를 입은 여성이 'I Love '라고 적힌 작은 그림 액자를 들고 있습니다. 조명이 그녀의 예술적 감각을 돋보이게 합니다.",
"가을 숲 속에서 두툼한 코트와 스카프를 두른 남성이 'I Love '라고 적힌 나뭇잎을 들고 있습니다. 차가운 가을의 분위기가 느껴집니다.",
"고풍스러운 서재에 앉아 있는 남성이 빈티지한 양복을 입고 있으며, 'I Love '라고 적힌 낡은 책을 들고 있습니다. 은은한 조명이 고급스러운 느낌을 더해줍니다.",
"별빛 가득한 밤하늘 아래 파란색 드레스를 입고 앉아 있는 여성이 'I Love '라고 적힌 별 모양 소품을 들고 있습니다. 몽환적이고 신비로운 분위기를 자아냅니다.",    
]    

css = """
footer {visibility: hidden;}
.container {max-width: 1200px; margin: auto; padding: 20px;}
.generate-box, .image-box {
    background-color: #f0f0f0;
    border-radius: 10px;
    padding: 20px;
    margin-bottom: 20px;
    height: 600px;  /* 고정된 높이 설정 */
    display: flex;
    flex-direction: column;
}
.image-box img {
    max-height: 100%;
    width: 100%;
    object-fit: contain;
}
.generate-box .row {display: flex; align-items: center; margin-bottom: 10px;}
.generate-box .row > * {margin-right: 10px;}
.generate-box .row > *:last-child {margin-right: 0;}
.advanced-options {background-color: #e0e0e0; border-radius: 10px; padding: 20px; margin-top: 20px;}
.examples-gallery {margin-top: 30px;}
"""

# --- Gradio UI ---
with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
    with gr.Column(elem_id="container"):
        gr.Markdown("# Open FLUX 1.1 Pro")
        gr.Markdown("Flux Schnell-based with no commercial restrictions, 4-step fast image generation with quality enhancement, and improved memory efficiency (VAE).")        

        with gr.Row():
            with gr.Column(scale=2):
                result = gr.Image(label="Generated Image", show_label=False, interactive=False, elem_classes="image-box")
            with gr.Column(scale=1):
                with gr.Column(elem_classes="generate-box"):
                    prompt = gr.Text(
                        label="Prompt (한글 입력 가능)",
                        placeholder="섹시한 여자와 남자, 속옷, 전신, 일요일",
                        lines=3,
                    )
                    generateBtn = gr.Button("Generate Image", variant="primary")

                    with gr.Column(elem_classes="advanced-options"):
                        with gr.Row():
                            seed = gr.Number(label="Seed", value=42)
                            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=DEFAULT_WIDTH)
                            height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)

        with gr.Column(elem_classes="examples-gallery"):
            gr.Markdown("### Gallery")
            gr.Examples(
                examples=examples,
                fn=generate_image,
                inputs=[prompt],
                outputs=[result, seed],
                cache_examples="lazy" 
            )

    generateBtn.click(
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed],
        outputs=[result, seed],
        show_progress="full",
        api_name="GenerateImage",
    )

# Launch the app
demo.launch()