Update app.py

app.py

@@ -1,112 +1,138 @@
+##############################################################################
+#                                app.py                                      #
+##############################################################################
 model_repo_id = "Freepik/F-Lite-Texture"
-model_name = "F Lite Texture" 
+model_name    = "F Lite Texture"
 
-from dotenv import load_dotenv
+from dotenv   import load_dotenv
 import gradio as gr
-import numpy as np
-import random
-import os
-import logging
-import google.generativeai as genai
-
-import spaces
-import torch
+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
 
-# Trick required because it is not a native diffusers model
+# ────────────────────────────────────────────────────────────────────────────
+# 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"]
 
-from f_lite.pipeline import APGConfig
-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)
 
-# Initialize Gemini API if API key is available
+# Gemini API 준비 (있을 때만 사용)
+gemini_available = False
 if os.getenv("GEMINI_API_KEY"):
-    gemini_available = True
     genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+    gemini_available = True
 else:
-    gemini_available = False
-    logging.warning("GEMINI_API_KEY not found in environment variables. Prompt enrichment will not work.")
+    logging.warning("GEMINI_API_KEY not found – prompt enrichment disabled.")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-if torch.cuda.is_available():
-    torch_dtype = torch.bfloat16
-else:
-    torch_dtype = torch.float32
+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()
+pipe.vae.enable_slicing(); pipe.vae.enable_tiling()
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1600
+# ────────────────────────────────────────────────────────────────────────────
+# 기본 값 설정
+# ────────────────────────────────────────────────────────────────────────────
+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"}
-  ]
+    "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 category, resolutions in RESOLUTIONS.items():
-    resolution_options.append([f"{category.capitalize()}", None])  # 카테고리 헤더용
-    for res in resolutions:
-        resolution_options.append([f"  {res['label']}", f"{category}:{res['width']}:{res['height']}"])
-
-def enrich_prompt_with_gemini(prompt, max_tokens=1024):
+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 os.getenv("GEMINI_API_KEY"):
-            return None, "GEMINI_API_KEY not found in environment variables."
-        
-        model = genai.GenerativeModel('gemini-1.5-flash')
-        enrichment_prompt = f"""
-        You are a prompt enhancer for image generation. 
-        Take the following basic prompt and make it longer, very descriptive, and detailed. 
-        Write the description in a paragraph, avoiding bullet points.
-        
-        Original prompt: {prompt}
-        
-        Enhanced prompt:
-        """
-        response = model.generate_content(enrichment_prompt, generation_config={
-            "max_output_tokens": max_tokens,
-            "temperature": 1,
-        })
-        enriched_prompt = response.text.strip()
-        return enriched_prompt, None
+        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:
-        error_message = f"Error enriching prompt: {str(e)}"
-        logging.error(error_message)
-        return None, error_message
-
-def update_resolution(resolution_value):
-    if not resolution_value:
+        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:
-        category, width, height = resolution_value.split(":")
-        return int(width), int(height)
-    except:
+        _, 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,
@@ -121,15 +147,14 @@ def infer(
     enable_apg,
     progress=gr.Progress(track_tqdm=True),
 ):
-    enriched_prompt_str = None
-    error_message_str = None
     generation_prompt = prompt
-    
+    enriched_prompt, enrich_err = None, None
+
     if use_prompt_enrichment and gemini_available:
-        enriched_prompt_str, error_message_str = enrich_prompt_with_gemini(prompt)
-        if enriched_prompt_str:
-            generation_prompt = enriched_prompt_str
-    
+        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)
 
@@ -140,224 +165,122 @@ def infer(
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
+        width=width, height=height,
         generator=generator,
-        apg_config=APGConfig(enabled=enable_apg)
+        apg_config=APGConfig(enabled=enable_apg),
     ).images[0]
-    
-    enriched_prompt_display_update = gr.update(visible=False)
-    enriched_prompt_text_update = gr.update(value="")
-    enrichment_error_update = gr.update(visible=False, value="")
 
-    if enriched_prompt_str:
-        enriched_prompt_display_update = gr.update(visible=True)
-        enriched_prompt_text_update = gr.update(value=enriched_prompt_str)
-    elif error_message_str:
-        enriched_prompt_display_update = gr.update(visible=True)
-        enrichment_error_update = gr.update(visible=True, value=error_message_str)
+    # UI 업데이트 제어
+    show_acc   = gr.update(visible=False)
+    show_text  = gr.update(value="")
+    show_error = gr.update(visible=False, value="")
 
-    return image, seed, enriched_prompt_display_update, enriched_prompt_text_update, enrichment_error_update
-
-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 pine branches frame 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 ambient teal light", 
-        None
-    ],
-    [
-        "A majestic white Arabian horse galloping across a shallow mirror lake at sunset, water droplets frozen mid-air, warm rim light outlining powerful muscles, captured with a 1/4000 s shutter and sweeping motion blur background", 
-        None
-    ],
-    [
-        "An ancient library carved inside a giant redwood tree, spiral staircases of polished root wood, glowing fireflies as reading lamps, shafts of emerald light pierce stained-leaf windows, ultra-detailed fantasy matte-painting", 
-        None
-    ],
-    [
-        "A haute-couture fashion portrait: model wearing a dress made entirely of iridescent butterfly wings, dramatic chiaroscuro lighting, deep-blue velvet backdrop, captured on Hasselblad with razor-sharp focus on eyes, 120 MP clarity", 
-        None
-    ],
-    [
-        "Cyberpunk samurai duo standing beneath a torrential neon rainstorm, reflective katana blades crackling with violet energy, holographic kanji glyphs drifting in air, cinematic anamorphic lens flares, dark gritty atmosphere", 
-        None
-    ],
-    [
-        "A whimsical steampunk airship festival above a Victorian harbor: brass dirigibles adorned with floral patterns, fireworks burst into cog-shaped sparks, golden hour light bathes billowing clouds, painterly style reminiscent of Hayao Miyazaki", 
-        None
-    ],
-]
+    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;
-}
+#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")
-        
+        gr.Markdown(f"# {model_name} Text-to-Image Demo")
+
+        # ── 입력 파트 ──
         with gr.Row(elem_classes="prompt-row"):
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-                scale=6
+            prompt_box = gr.Text(
+                value=DEFAULT_PROMPT,              # 기본 프롬프트
+                label="Prompt", show_label=False,
+                max_lines=1, placeholder="Enter your prompt",
+                container=False, scale=6,
             )
-            # 1) Enrich 기본값 True / 표시 안 함(visible=False)
             use_prompt_enrichment = gr.Checkbox(
-                label="Enrich", 
-                value=True,    # default True
-                visible=False  # Hide from UI
+                label="Enrich", value=True, visible=False  # 숨김 + 기본 True
             )
             run_button = gr.Button("Run", scale=1, variant="primary", min_width=100)
-        
-        result = gr.Image(label="Result", show_label=False)
-        
-        enriched_prompt_display = gr.Accordion("Enriched Prompt", open=False, visible=False)
-        with enriched_prompt_display:
-            enriched_prompt_text = gr.Textbox(
-                label="Enriched Prompt",
-                interactive=False,
-                lines=8
-            )
-            enrichment_error = gr.Textbox(
-                label="Error",
-                visible=False,
-                interactive=False,
-            )
-        
+
+        # ── 결과 이미지 ── (초기값: 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",
-                visible=True,
+                label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt"
             )
-            
-            with gr.Tabs() as resolution_tabs:
+            with gr.Tabs():
                 with gr.TabItem("Preset Resolutions"):
-                    resolution_dropdown = gr.Dropdown(
+                    resolution_dd = gr.Dropdown(
                         label="Resolution",
                         choices=resolution_options,
-                        # 2) 디폴트로 1600x896(수평 최대 크기)
-                        value="horizontal:1600:896",
+                        value="horizontal:1600:896",   # 가장 큰 preset
                         type="value"
                     )
-                
                 with gr.TabItem("Custom Resolution"):
                     with gr.Row():
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            # set default to 1600
-                            value=1600,
+                        width_sl = gr.Slider(
+                            label="Width", minimum=256, maximum=MAX_IMAGE_SIZE,
+                            step=32, value=1600
                         )
-                        height = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            # set default to 896
-                            value=896,
+                        height_sl = gr.Slider(
+                            label="Height", minimum=256, maximum=MAX_IMAGE_SIZE,
+                            step=32, value=896
                         )
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=42,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+            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_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=15.0,
-                    step=0.1,
-                    value=6, # Keep default
-                )
-                enable_apg = gr.Checkbox(
-                    label="Enable APG",
-                    value=True,
-                )
-
-                # 3) num_inference_steps를 50으로 (최대값)
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=50,  # "highest quality" default
-                )
-
-        # 제한된 길이
-        max_length = 180
-
-        def set_example_and_disable_enrichment(example, current_checkbox_value):
-            return example, gr.update(value=False)
-        
+                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=50)
+
+        # ── 예제 프롬프트 목록 ──
+        def set_example(example, _):
+            return example, gr.update(value=False)  # enrichment 끄기
+        max_len = 180
         gr.Examples(
             examples=examples,
-            inputs=[prompt, use_prompt_enrichment],
-            outputs=[prompt, use_prompt_enrichment],
-            fn=set_example_and_disable_enrichment,
-            example_labels=[ex[0][:max_length] + "..." if len(ex[0]) > max_length else ex[0] for ex in 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_dropdown.change(
-        fn=update_resolution,
-        inputs=resolution_dropdown,
-        outputs=[width, height]
-    )
-    
+    # ── 상호작용 연결 ──
+    resolution_dd.change(fn=update_resolution, inputs=resolution_dd, outputs=[width_sl, height_sl])
+
     gr.on(
-        triggers=[run_button.click, prompt.submit],
+        triggers=[run_button.click, prompt_box.submit],
         fn=infer,
         inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            use_prompt_enrichment,
-            enable_apg,
+            prompt_box, negative_prompt, seed_sl, random_seed,
+            width_sl, height_sl, guidance_sl, steps_sl,
+            use_prompt_enrichment, enable_apg
         ],
-        outputs=[result, seed, enriched_prompt_display, enriched_prompt_text, enrichment_error],
+        outputs=[result_img, seed_sl, enrich_acc, enrich_txt, enrich_error],
     )
-    
 
+# ────────────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     demo.launch()