app.py

##############################################################################
#                                app.py                                      #
##############################################################################
model_repo_id = "Freepik/F-Lite-Texture"
model_name    = "F Lite Texture"

from dotenv   import load_dotenv
import gradio as gr
import numpy   as np
import random, os, logging, google.generativeai as genai, spaces, torch

from f_lite import FLitePipeline
from f_lite.pipeline import APGConfig

# ────────────────────────────────────────────────────────────────────────────
# diffusers helper (필수 트릭 – DiT 클래스를 diffusers로 인식시키기 위함)
# ────────────────────────────────────────────────────────────────────────────
from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES
LOADABLE_CLASSES["f_lite"]            = LOADABLE_CLASSES["f_lite.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
ALL_IMPORTABLE_CLASSES["DiT"]         = ["save_pretrained", "from_pretrained"]

# ────────────────────────────────────────────────────────────────────────────
# 환경 설정 / 모델 로드
# ────────────────────────────────────────────────────────────────────────────
load_dotenv()
logging.basicConfig(level=logging.INFO)

# Gemini API 준비 (있을 때만 사용)
gemini_available = False
if os.getenv("GEMINI_API_KEY"):
    genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
    gemini_available = True
else:
    logging.warning("GEMINI_API_KEY not found – prompt enrichment disabled.")

device      = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

pipe = FLitePipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe.to(device)
pipe.vae.enable_slicing(); pipe.vae.enable_tiling()

# ────────────────────────────────────────────────────────────────────────────
# 기본 값 설정
# ────────────────────────────────────────────────────────────────────────────
MAX_SEED        = np.iinfo(np.int32).max
MAX_IMAGE_SIZE  = 1600

RESOLUTIONS = {
    "horizontal": [
        {"width": 1344, "height": 896, "label": "1344×896"},
        {"width": 1152, "height": 768, "label": "1152×768"},
        {"width": 960 , "height": 640, "label": "960×640"},
        {"width": 1600, "height": 896, "label": "1600×896"}
    ],
    "vertical": [
        {"width": 896 , "height": 1344, "label": "896×1344"},
        {"width": 768 , "height": 1152, "label": "768×1152"},
        {"width": 640 , "height": 960 , "label": "640×960"},
        {"width": 896 , "height": 1600, "label": "896×1600"}
    ],
    "square": [
        {"width": 1216, "height": 1216, "label": "1216×1216"},
        {"width": 1024, "height": 1024, "label": "1024×1024"}
    ]
}
DEFAULT_RESOLUTION = {"width": 1024, "height": 1024, "label": "1024×1024"}

# ────────────────────────────────────────────────────────────────────────────
# 해상도 드롭다운 옵션 생성
# ────────────────────────────────────────────────────────────────────────────
resolution_options = []
for cat, res_list in RESOLUTIONS.items():
    resolution_options.append([f"{cat.capitalize()}", None])
    for r in res_list:
        resolution_options.append([f"  {r['label']}", f"{cat}:{r['width']}:{r['height']}"])

# ────────────────────────────────────────────────────────────────────────────
# Prompt enrichment (Gemini)
# ────────────────────────────────────────────────────────────────────────────
def enrich_prompt_with_gemini(prompt: str, max_tokens: int = 1024):
    """Gemini-based prompt expansion (에러 시 원본 유지)."""
    try:
        if not gemini_available:
            return None, "Gemini unavailable."
        model = genai.GenerativeModel("gemini-1.5-flash")
        ask   = (
            "You are an exceptional prompt enhancer for text-to-image generation.\n"
            "Rewrite the following prompt so it becomes richly detailed, cinematic, and vivid.\n"
            "Return ONE descriptive paragraph only.\n\n"
            f"Original prompt: {prompt}\n\nEnhanced prompt:"
        )
        out = model.generate_content(
            ask,
            generation_config={"max_output_tokens": max_tokens, "temperature": 1},
        )
        return out.text.strip(), None
    except Exception as e:
        logging.error(f"Gemini error: {e}")
        return None, f"Gemini error: {e}"

# ────────────────────────────────────────────────────────────────────────────
# 해상도 업데이트
# ────────────────────────────────────────────────────────────────────────────
def update_resolution(sel: str):
    if not sel:
        return DEFAULT_RESOLUTION["width"], DEFAULT_RESOLUTION["height"]
    try:
        _, w, h = sel.split(":")
        return int(w), int(h)
    except ValueError:
        return DEFAULT_RESOLUTION["width"], DEFAULT_RESOLUTION["height"]

# ────────────────────────────────────────────────────────────────────────────
# 예시 프롬프트 10개
# ────────────────────────────────────────────────────────────────────────────
examples = [
    ["An ultra-detailed macro photograph of a dew-covered rainbow beetle perched on a spiralling fern unfurling at dawn, back-lit by golden sunrise, bokeh background, 200-mm lens, f/2.8, vivid colors, cinematic lighting", None],
    ["A retro-futuristic cityscape at night inspired by Syd Mead: neon-drenched streets reflect glistening rain, flying cars leave light-trail arcs between towering holographic billboards, 35 mm film grain, wide-angle perspective", None],
    ["An elegant 18th-century ballroom rendered in photorealistic 8K, crystal chandeliers scattering prismatic light across polished marble floors, dancers in flowing silk gowns twirl mid-motion, captured like a long-exposure still", None],
    ["A serene Japanese onsen nestled in a snowy mountain valley, steam rising into crisp twilight air, red lanterns glowing softly, snow-laden pines framing the scene, shot on medium-format analog with natural film tones", None],
    ["Hyper-real illustration of an astronaut in a translucent spacesuit tending a floating bonsai inside a zero-g greenhouse aboard an orbital station, earthrise through panoramic windows, bioluminescent plants provide teal ambience", None],
    ["A majestic white Arabian horse galloping across a mirror-like lake at sunset, droplets frozen mid-air, warm rim light outlining powerful muscles, captured at 1/4000 s with sweeping motion-blur background", None],
    ["An ancient library hollowed inside a colossal redwood tree, spiral root staircases, glowing fireflies as lamps, shafts of emerald light pierce stained-leaf windows, ultra-detailed fantasy matte painting", None],
    ["A haute-couture portrait: model in a gown composed of iridescent butterfly wings, dramatic chiaroscuro, deep-blue velvet backdrop, shot on Hasselblad with razor-sharp eye focus, 120 MP clarity", None],
    ["Cyberpunk samurai duo beneath torrential neon rain, reflective katanas crackling with violet energy, holographic kanji drifting, cinematic anamorphic lens flares, gritty atmosphere", None],
    ["A whimsical steampunk airship festival above a Victorian harbor: brass dirigibles with floral patterns, cog-shaped fireworks burst at golden hour, painterly style reminiscent of Miyazaki", None],
]

# 첫 예시 프롬프트를 기본값으로 사용
DEFAULT_PROMPT = examples[0][0]

# ────────────────────────────────────────────────────────────────────────────
# 추론 함수
# ────────────────────────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    use_prompt_enrichment,
    enable_apg,
    progress=gr.Progress(track_tqdm=True),
):
    generation_prompt = prompt
    enriched_prompt, enrich_err = None, None

    if use_prompt_enrichment and gemini_available:
        enriched_prompt, enrich_err = enrich_prompt_with_gemini(prompt)
        if enriched_prompt:
            generation_prompt = enriched_prompt

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=generation_prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width, height=height,
        generator=generator,
        apg_config=APGConfig(enabled=enable_apg),
    ).images[0]

    # UI 업데이트 제어
    show_acc   = gr.update(visible=False)
    show_text  = gr.update(value="")
    show_error = gr.update(visible=False, value="")

    if enriched_prompt:
        show_acc  = gr.update(visible=True)
        show_text = gr.update(value=enriched_prompt)
    elif enrich_err:
        show_acc   = gr.update(visible=True)
        show_error = gr.update(visible=True, value=enrich_err)

    return image, seed, show_acc, show_text, show_error

# ────────────────────────────────────────────────────────────────────────────
# Gradio UI
# ────────────────────────────────────────────────────────────────────────────
css = """
#col-container {margin:0 auto; max-width:1024px;}
.prompt-row > .gr-form{gap:0.5rem !important; align-items:center;}
"""

with gr.Blocks(css=css, theme="ParityError/Interstellar") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"# {model_name} Text-to-Image Demo")

        # ── 입력 파트 ──
        with gr.Row(elem_classes="prompt-row"):
            prompt_box = gr.Text(
                value=DEFAULT_PROMPT,              # 기본 프롬프트
                label="Prompt", show_label=False,
                max_lines=1, placeholder="Enter your prompt",
                container=False, scale=6,
            )
            use_prompt_enrichment = gr.Checkbox(
                label="Enrich", value=True, visible=False  # 숨김 + 기본 True
            )
            run_button = gr.Button("Run", scale=1, variant="primary", min_width=100)

        # ── 결과 이미지 ── (초기값: image1.webp)
        result_img = gr.Image(
            value="image1.webp",                 # 같은 경로 이미지 표시
            label="Result", show_label=False
        )

        # Enriched prompt 표시용 아코디언
        enrich_acc = gr.Accordion("Enriched Prompt", open=False, visible=False)
        with enrich_acc:
            enrich_txt   = gr.Textbox(label="Enriched Prompt", interactive=False, lines=8)
            enrich_error = gr.Textbox(label="Error", visible=False, interactive=False)

        # ── 고급 설정 ──
        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt"
            )
            with gr.Tabs():
                with gr.TabItem("Preset Resolutions"):
                    resolution_dd = gr.Dropdown(
                        label="Resolution",
                        choices=resolution_options,
                        value="horizontal:1600:896",   # 가장 큰 preset
                        type="value"
                    )
                with gr.TabItem("Custom Resolution"):
                    with gr.Row():
                        width_sl = gr.Slider(
                            label="Width", minimum=256, maximum=MAX_IMAGE_SIZE,
                            step=32, value=1600
                        )
                        height_sl = gr.Slider(
                            label="Height", minimum=256, maximum=MAX_IMAGE_SIZE,
                            step=32, value=896
                        )

            seed_sl = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
            random_seed = gr.Checkbox(label="Randomize seed", value=False)

            with gr.Row():
                guidance_sl = gr.Slider(label="Guidance scale", minimum=0, maximum=15, step=0.1, value=6)
                enable_apg  = gr.Checkbox(label="Enable APG", value=True)
                steps_sl    = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=30)

        # ── 예제 프롬프트 목록 ──
        def set_example(example, _):
            return example, gr.update(value=False)  # enrichment 끄기
        max_len = 180
        gr.Examples(
            examples=examples,
            inputs=[prompt_box, use_prompt_enrichment],
            outputs=[prompt_box, use_prompt_enrichment],
            fn=set_example,
            example_labels=[ex[0][:max_len] + "..." if len(ex[0]) > max_len else ex[0] for ex in examples]
        )

        gr.Markdown(f"[{model_name} Model Card and Weights](https://huggingface.co/{model_repo_id})")

    # ── 상호작용 연결 ──
    resolution_dd.change(fn=update_resolution, inputs=resolution_dd, outputs=[width_sl, height_sl])

    gr.on(
        triggers=[run_button.click, prompt_box.submit],
        fn=infer,
        inputs=[
            prompt_box, negative_prompt, seed_sl, random_seed,
            width_sl, height_sl, guidance_sl, steps_sl,
            use_prompt_enrichment, enable_apg
        ],
        outputs=[result_img, seed_sl, enrich_acc, enrich_txt, enrich_error],
    )

# ────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    demo.launch()