############################################################################## # 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()