Update app.py

app.py

@@ -1,14 +1,62 @@
+"""
+FLUX.1 Kontext Style Transfer
+==============================
+Updated: 2025‑07‑12
+---------------------------------
+이 스크립트는 Hugging Face **FLUX.1‑Kontext‑dev** 모델과
+22 종의 스타일 LoRA 가중치를 이용해 이미지를 다양한 예술
+스타일로 변환하는 Gradio 데모입니다.
+
+주요 개선 사항
+--------------
+1. **모델 캐싱** – `snapshot_download()`로 실행 시작 시 한 번만
+   모델과 LoRA를 캐싱해 이후 GPU 잡에서도 재다운로드가 없도록
+   함.
+2. **GPU VRAM 자동 판별** – GPU VRAM이 24 GB 미만이면
+   `torch.float16` / `enable_sequential_cpu_offload()`를 자동 적용.
+3. **단일 로딩 메시지** – Gradio `gr.Info()` 메시지가 최초 1회만
+   표시되도록 수정.
+4. **버그 픽스** – seed 처리, LoRA 언로드, 이미지 리사이즈 로직
+   등 세부 오류 수정.
+
+------------------------------------------------------------
+"""
+import os
 import gradio as gr
 import spaces
 import torch
+from huggingface_hub import snapshot_download
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 from PIL import Image
-import os
 
-# ----------------------------------------------
-# Style → LoRA file mapping
-# ----------------------------------------------
+# ------------------------------------------------------------------
+# 환경 설정 & 모델 / LoRA 사전 다운로드
+# ------------------------------------------------------------------
+# 큰 파일을 빠르게 받도록 가속 플래그 활성화
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+MODEL_ID = "black-forest-labs/FLUX.1-Kontext-dev"
+LORA_REPO = "Owen777/Kontext-Style-Loras"
+CACHE_DIR = os.getenv("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+
+# --- 최초 실행 시에만 다운로드(이미 캐시에 있으면 건너뜀) ---
+MODEL_DIR = snapshot_download(
+    repo_id=MODEL_ID,
+    cache_dir=CACHE_DIR,
+    resume_download=True,
+    token=True  # HF 토큰(필요 시 환경변수 HF_TOKEN 지정)
+)
+LORA_DIR = snapshot_download(
+    repo_id=LORA_REPO,
+    cache_dir=CACHE_DIR,
+    resume_download=True,
+    token=True
+)
+
+# ------------------------------------------------------------------
+# 스타일 → LoRA 파일 매핑 & 설명
+# ------------------------------------------------------------------
 STYLE_LORA_MAP = {
     "3D_Chibi": "3D_Chibi_lora_weights.safetensors",
     "American_Cartoon": "American_Cartoon_lora_weights.safetensors",
@@ -31,116 +79,106 @@ STYLE_LORA_MAP = {
     "Vector": "Vector_lora_weights.safetensors",
     "Picasso": "Picasso_lora_weights.safetensors",
     "Macaron": "Macaron_lora_weights.safetensors",
-    "Rick_Morty": "Rick_Morty_lora_weights.safetensors"
+    "Rick_Morty": "Rick_Morty_lora_weights.safetensors",
 }
 
-# ----------------------------------------------
-# Style descriptions (툴팁용)
-# ----------------------------------------------
 STYLE_DESCRIPTIONS = {
-    "3D_Chibi": "귀여운 SD 캐릭터의 3D 느낌",
-    "American_Cartoon": "고전적인 미국 카툰 스타일",
-    "Chinese_Ink": "수묵화의 번짐과 농담 표현",
-    "Clay_Toy": "찰흙·플라스틴 장난감質",
-    "Fabric": "섬유·패브릭 질감",
-    "Ghibli": "지브리풍 따뜻한 색감 & 연필선",
-    "Irasutoya": "일러스토야 미니멀 평면 그림",
-    "Jojo": "죠죠의 기묘한 모험 망가 터치",
-    "Oil_Painting": "유화 붓터치와 질감",
-    "Pixel": "16/32‑bit 레트로 픽셀아트",
-    "Snoopy": "피너츠 스트립(스누피) 스타일",
-    "Poly": "로우폴리 3D 기하학적 스타일",
-    "LEGO": "레고 블록 조립 스타일",
-    "Origami": "종이접기 질감·각도",
-    "Pop_Art": "팝아트의 선명한 색감과 도트",
-    "Van_Gogh": "반 고흐의 굵은 임파스토",
-    "Paper_Cutting": "종이 오리기 실루엣",
-    "Line": "깔끔한 라인 드로잉",
-    "Vector": "벡터 그래픽·플랫 디자인",
-    "Picasso": "피카소식 입체주의 큐비즘",
-    "Macaron": "파스텔 마카롱톤 부드러움",
-    "Rick_Morty": "릭 앤 모티 애니메이션 스타일"
+    "3D_Chibi": "Cute, miniature 3D character style with big heads",
+    "American_Cartoon": "Classic American animation style",
+    "Chinese_Ink": "Traditional Chinese ink painting aesthetic",
+    "Clay_Toy": "Playful clay/plasticine toy appearance",
+    "Fabric": "Soft, textile-like rendering",
+    "Ghibli": "Studio Ghibli's distinctive anime style",
+    "Irasutoya": "Simple, flat Japanese illustration style",
+    "Jojo": "JoJo's Bizarre Adventure manga style",
+    "Oil_Painting": "Classic oil painting texture and strokes",
+    "Pixel": "Retro pixel art style",
+    "Snoopy": "Peanuts comic strip style",
+    "Poly": "Low-poly 3D geometric style",
+    "LEGO": "LEGO brick construction style",
+    "Origami": "Paper folding art style",
+    "Pop_Art": "Bold, colorful pop art style",
+    "Van_Gogh": "Van Gogh's expressive brushstroke style",
+    "Paper_Cutting": "Paper cut-out art style",
+    "Line": "Clean line art/sketch style",
+    "Vector": "Clean vector graphics style",
+    "Picasso": "Cubist art style inspired by Picasso",
+    "Macaron": "Soft, pastel macaron-like style",
+    "Rick_Morty": "Rick and Morty cartoon style",
 }
 
-# 전역 파이프라인
-pipe = None
-
-
-def get_dtype():
-    """최적 dtype(bf16 지원 시 bf16, 아니면 fp16) 선택"""
-    if torch.cuda.is_available():
-        major, _ = torch.cuda.get_device_capability()
-        if major >= 8:  # Ada/Hopper GPU는 bf16 본격 지원
-            return torch.bfloat16
-    return torch.float16
+# ------------------------------------------------------------------
+# 파이프라인 로더 (단일 인스턴스)
+# ------------------------------------------------------------------
+_pipeline = None  # 내부 글로벌 캐시
 
 
 def load_pipeline():
-    """FluxKontext 파이프라인을 (지연)로드"""
-    global pipe
-    if pipe is None:
-        gr.Info("⬇️ FLUX.1‑Kontext 모델을 다운로드합니다…")
-
-        pipe = FluxKontextPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-Kontext-dev",
-            torch_dtype=get_dtype(),
-            resume_download=True,
-        )
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        pipe.to(device)
-
-        # VRAM 절약 모드
-        if torch.cuda.is_available():
-            pipe.enable_sequential_cpu_offload()
-        pipe.vae.enable_tiling()
+    """Load (or return cached) FluxKontextPipeline."""
+    global _pipeline
+    if _pipeline is not None:
+        return _pipeline
+
+    # VRAM이 24 GB 미만이면 FP16 사용 + CPU 오프로딩
+    dtype = torch.bfloat16
+    vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
+    if vram_gb < 24:
+        dtype = torch.float16
+
+    gr.Info("FLUX.1‑Kontext 파이프라인 로딩 중… (최초 1회)")
+
+    pipe = FluxKontextPipeline.from_pretrained(
+        MODEL_DIR,
+        torch_dtype=dtype,
+        local_files_only=True,
+    )
 
-    return pipe
+    pipe.to("cuda")
 
+    if vram_gb < 24:
+        pipe.enable_sequential_cpu_offload()
+    else:
+        pipe.enable_model_cpu_offload()
 
-@spaces.GPU(duration=600)  # 초기 다운로드 시간 확보
-def style_transfer(
-    input_image,
-    style_name,
-    prompt_suffix,
-    num_inference_steps,
-    guidance_scale,
-    seed,
-):
-    """선택한 스타일로 이미지 변환"""
+    _pipeline = pipe
+    return _pipeline
 
+# ------------------------------------------------------------------
+# 스타일 변환 함수 (Spaces GPU 잡)
+# ------------------------------------------------------------------
+@spaces.GPU(duration=600)
+def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps, guidance_scale, seed):
+    """Apply selected style to the uploaded image."""
     if input_image is None:
-        gr.Warning("🖼️ 먼저 이미지를 업로드하세요!")
+        gr.Warning("Please upload an image first!")
         return None
 
     try:
         pipe = load_pipeline()
 
-        # 시드 고정 (0이면 난수)
-        generator = None
-        if seed and int(seed) != 0:
-            generator = torch.Generator(device=pipe.device).manual_seed(int(seed))
+        # --- Torch Generator 설정 ---
+        if seed > 0:
+            generator = torch.Generator(device="cuda").manual_seed(int(seed))
+        else:
+            generator = None  # random
 
-        # 이미지 로드 & 리사이즈
-        img = load_image(input_image) if isinstance(input_image, str) else input_image
+        # --- 입력 이미지 전처리 ---
+        img = input_image if isinstance(input_image, Image.Image) else load_image(input_image)
         img = img.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS)
 
-        # LoRA 로딩
+        # --- LoRA 로드 ---
         lora_file = STYLE_LORA_MAP[style_name]
         adapter_name = "style"
-        pipe.load_lora_weights(
-            "Owen777/Kontext-Style-Loras",
-            weight_name=lora_file,
-            adapter_name=adapter_name,
-        )
-        pipe.set_adapters([adapter_name], adapter_weights=[1.0])
+        pipe.load_lora_weights(LORA_DIR, weight_name=lora_file, adapter_name=adapter_name)
+        pipe.set_adapters([adapter_name], [1.0])
 
-        # 프롬프트 구성
-        prompt = f"Turn this image into the {style_name.replace('_', ' ')} style."
+        # --- 프롬프트 빌드 ---
+        human_readable_style = style_name.replace("_", " ")
+        prompt = f"Turn this image into the {human_readable_style} style."
         if prompt_suffix and prompt_suffix.strip():
             prompt += f" {prompt_suffix.strip()}"
 
-        gr.Info("🎨 이미지를 생성 중…")
+        gr.Info("Generating styled image… (24‑60 s)")
 
         result = pipe(
             image=img,
@@ -152,111 +190,85 @@ def style_transfer(
             width=1024,
         )
 
-        # LoRA 언로드 및 캐시 정리
-        pipe.unload_lora_weights()
+        # --- LoRA 언로드 & GPU 메모리 해제 ---
+        pipe.unload_lora_weights(adapter_name=adapter_name)
         torch.cuda.empty_cache()
 
         return result.images[0]
 
     except Exception as e:
-        gr.Error(f"🚨 오류: {e}")
         torch.cuda.empty_cache()
+        gr.Error(f"Error during style transfer: {e}")
         return None
 
-
-def update_description(style):
-    return STYLE_DESCRIPTIONS.get(style, "")
-
-
+# ------------------------------------------------------------------
+# Gradio UI 정의
+# ------------------------------------------------------------------
 with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🎨 FLUX.1 Kontext Style Transfer
-        FLUX.1‑Kontext‑dev 모델과 22개의 고품질 LoRA로 이미지를 다양한 예술 스타일로 변환하세요.
-        """
-    )
+    gr.Markdown("""
+    # 🎨 FLUX.1 Kontext Style Transfer
+
+    업로드한 이미지를 22 종의 예술 스타일로 변환하세요!
+    (모델 / LoRA는 최초 실행 시에만 다운로드되며, 이후 실행은 빠릅니다.)
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(label="Upload Image", type="pil", height=400)
-
             style_dropdown = gr.Dropdown(
                 choices=list(STYLE_LORA_MAP.keys()),
                 value="Ghibli",
                 label="Select Style",
-                info="Choose from 22 different artistic styles",
-            )
-
-            style_info = gr.Textbox(
-                label="Style Description",
-                value=STYLE_DESCRIPTIONS["Ghibli"],
-                interactive=False,
-                lines=2,
-            )
-
-            prompt_suffix = gr.Textbox(
-                label="Additional Instructions (Optional)",
-                placeholder="예: 'make it more colorful', 'add dramatic lighting' …",
-                lines=2,
             )
+            style_info = gr.Textbox(label="Style Description", value=STYLE_DESCRIPTIONS["Ghibli"], interactive=False, lines=2)
+            prompt_suffix = gr.Textbox(label="Additional Instructions (Optional)", placeholder="e.g. add dramatic lighting", lines=2)
 
             with gr.Accordion("Advanced Settings", open=False):
-                num_steps = gr.Slider(
-                    10,
-                    50,
-                    value=24,
-                    step=1,
-                    label="Inference Steps",
-                    info="더 높을수록 품질↑ 속도↓",
-                )
-                guidance = gr.Slider(
-                    1.0,
-                    7.5,
-                    value=2.5,
-                    step=0.1,
-                    label="Guidance Scale",
-                    info="프롬프트 준수 정도",
-                )
-                seed = gr.Number(label="Seed (0 = Random)", value=0)
-
-            generate_btn = gr.Button(
-                "🎨 Transform Image", variant="primary", size="lg"
-            )
+                num_steps = gr.Slider(minimum=10, maximum=50, value=24, step=1, label="Inference Steps")
+                guidance = gr.Slider(minimum=1.0, maximum=7.5, value=2.5, step=0.1, label="Guidance Scale")
+                seed = gr.Number(label="Seed (0 = random)", value=42)
+
+            generate_btn = gr.Button("🎨 Transform Image", variant="primary", size="lg")
 
         with gr.Column(scale=1):
             output_image = gr.Image(label="Styled Result", type="pil", height=400)
-
-            gr.Markdown(
-                """### 💡 Tips:
-                - 모든 이미지는 1024×1024로 리사이즈됩니다.
-                - 첫 실행 시 7 GB 모델 다운로드가 필요합니다.
-                - 스타일 변환은 약 30‑60 초 소요됩니다.
-                - 다른 스타일도 시험해 보세요!"""
-            )
-
-    # 이벤트 바인딩
-    style_dropdown.change(update_description, [style_dropdown], [style_info])
-
-    generate_btn.click(
-        style_transfer,
-        inputs=[
-            input_image,
-            style_dropdown,
-            prompt_suffix,
-            num_steps,
-            guidance,
-            seed,
+            gr.Markdown("""
+            ### 💡 Tips
+            * 이미지 크기는 1024×1024로 리사이즈됩니다.
+            * 최초 1회 모델 + LoRA 다운로드 후에는 **캐시**를 사용하므로 10‑20 s 내 완료됩니다.
+            * "Additional Instructions"에 색감·조명·효과 등을 영어로 간단히 적으면 결과를 세밀하게 제어할 수 있습니다.
+            """)
+
+    # --- 스타일 설명 자동 업데이트 ---
+    def _update_desc(style):
+        return STYLE_DESCRIPTIONS.get(style, "")
+
+    style_dropdown.change(fn=_update_desc, inputs=[style_dropdown], outputs=[style_info])
+
+    # --- 예제 ---
+    gr.Examples(
+        examples=[
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "Ghibli", ""],
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "3D_Chibi", "make it extra cute"],
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "Van_Gogh", "with swirling sky"],
         ],
-        outputs=[output_image],
+        inputs=[input_image, style_dropdown, prompt_suffix],
+        outputs=output_image,
+        fn=lambda img, style, prompt: style_transfer(img, style, prompt, 24, 2.5, 42),
+        cache_examples=False,
     )
 
-    gr.Markdown(
-        """
-        ---
-        Created with ❤️ by [Black‑Forest Labs](https://huggingface.co/black-forest-labs) & 
-        [Owen777/Kontext‑Style‑Loras](https://huggingface.co/Owen777/Kontext-Style-Loras)
-        """
+    # --- 버튼 연결 ---
+    generate_btn.click(
+        fn=style_transfer,
+        inputs=[input_image, style_dropdown, prompt_suffix, num_steps, guidance, seed],
+        outputs=output_image,
     )
 
+    gr.Markdown("""
+    ---
+    **Created with ❤️ by GiniGEN (2025)**
+    """)
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(inline=False)