Update app.py

app.py

@@ -1,49 +1,35 @@
 """
-FLUX.1 Kontext Style Transfer
-==============================
-Updated: 2025‑07‑12 (HF_TOKEN 지원)
----------------------------------
-이 스크립트는 Hugging Face **FLUX.1‑Kontext‑dev** 모델과
-22 종의 스타일 LoRA 가중치를 이용해 이미지를 다양한 예술
-스타일로 변환하는 Gradio 데모입니다.
-
-주요 개선 사항
---------------
-1. **HF_TOKEN 환경변수 지원** – `LocalTokenNotFoundError`를
-   방지하기 위해 `HF_TOKEN`(또는 런타임 로그인) 값을 자동으로
-   감지해 `snapshot_download()`에 전달합니다.
-2. **모델 캐싱** – `snapshot_download()`로 실행 시작 시 한 번만
-   모델과 LoRA를 캐싱.
-3. **GPU VRAM 자동 판별** – 24 GB 미만이면 FP16 / CPU offload.
-4. **단일 로딩 메시지** – Gradio `gr.Info()` 메시지가 최초 1회만
-   표시되도록 유지.
-5. **버그 픽스** – seed 처리, LoRA 언로드, 이미지 리사이즈.
-
-------------------------------------------------------------
+FLUX.1 Kontext Style Transfer
+=============================
+Updated : 2025‑07‑12 (HF_TOKEN 지원 + 전체 코드 완성)
+---------------------------------------------------
+Gradio 데모로 이미지를 22 종 예술 스타일로 변환합니다.
+- **HF_TOKEN** 환경변수를 자동 인식해 라이선스 모델 다운로드 오류를 방지합니다.
+- 최초 실행 시 모델·LoRA를 캐시에 받아 두고, 이후에는 재다운로드가 없습니다.
+- GPU VRAM을 감지하여 24 GB 미만에서는 FP16 + CPU offload를 사용합니다.
+- 파이프라인·LoRA 로딩 메시지는 최초 1회만 표시됩니다.
 """
+
 import os
 import gradio as gr
 import spaces
 import torch
-from huggingface_hub import snapshot_download, login as hf_login
+from huggingface_hub import snapshot_download
 from huggingface_hub.errors import LocalTokenNotFoundError
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 from PIL import Image
 
 # ------------------------------------------------------------------
-# 환경 설정 & 모델 / LoRA 사전 다운로드
+# 환경 설정 & 모델 / LoRA 사전 다운로드
 # ------------------------------------------------------------------
-# 큰 파일을 빠르게 받도록 가속 플래그 활성화
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+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"))
+HF_TOKEN = os.getenv("HF_TOKEN")  # 런타임에 주입하거나 Secrets 사용
 
-# --- HF 토큰 처리 --------------------------------------------------
-HF_TOKEN = os.getenv("HF_TOKEN")  # 런타임에 환경변수로 주입하거나
-# docker / Space 설정의 Secrets → HF_TOKEN 로 등록
 
 def _download_with_token(repo_id: str) -> str:
     """Download repo snapshot with optional token handling."""
@@ -52,16 +38,16 @@ def _download_with_token(repo_id: str) -> str:
             repo_id=repo_id,
             cache_dir=CACHE_DIR,
             resume_download=True,
-            token=HF_TOKEN if HF_TOKEN else True,  # True → HF_CACHED_TOKEN(or login)
+            token=HF_TOKEN if HF_TOKEN else True,  # True → 로그인 세션 사용
         )
     except LocalTokenNotFoundError:
-        # 미로그인 + 필수 동의 모델이면 에러 메시지 출력 후 종료
         raise RuntimeError(
             "Hugging Face 토큰이 필요합니다. 환경변수 HF_TOKEN을 설정하거나\n"
             "`huggingface-cli login`으로 로그인해 주세요."
         )
 
-# --- 최초 실행 시에만 다운로드(이미 캐시에 있으면 건너뜀) ---
+
+# 최초 실행 시 캐시에만 다운로드 (이미 있으면 즉시 반환)
 MODEL_DIR = _download_with_token(MODEL_ID)
 LORA_DIR = _download_with_token(LORA_REPO)
 
@@ -119,21 +105,20 @@ STYLE_DESCRIPTIONS = {
 }
 
 # ------------------------------------------------------------------
-# 파이프라인 로더 (단일 인스턴스)
+# 파이프라인 로더 (싱글턴)
 # ------------------------------------------------------------------
-_pipeline = None  # 내부 글로벌 캐시
+_pipeline = None
+
 
 def load_pipeline():
-    """Load (or return cached) FluxKontextPipeline."""
+    """Load or return cached FluxKontextPipeline."""
     global _pipeline
     if _pipeline is not None:
         return _pipeline
 
-    # VRAM이 24 GB 미만이면 FP16 사용 + CPU 오프로딩
-    dtype = torch.bfloat16
+    # VRAM 판별 → dtype & offload 설정
     vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
-    if vram_gb < 24:
-        dtype = torch.float16
+    dtype = torch.bfloat16 if vram_gb >= 24 else torch.float16
 
     gr.Info("FLUX.1‑Kontext 파이프라인 로딩 중… (최초 1회)")
 
@@ -141,9 +126,7 @@ def load_pipeline():
         MODEL_DIR,
         torch_dtype=dtype,
         local_files_only=True,
-    )
-
-    pipe.to("cuda")
+    ).to("cuda")
 
     if vram_gb < 24:
         pipe.enable_sequential_cpu_offload()
@@ -154,9 +137,10 @@ def load_pipeline():
     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:
@@ -166,28 +150,28 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
     try:
         pipe = load_pipeline()
 
-        # --- Torch Generator 설정 ---
+        # Torch Generator (seed 고정 시 재현 가능)
         generator = None
         if seed and int(seed) > 0:
             generator = torch.Generator(device="cuda").manual_seed(int(seed))
 
-        # --- 입력 이미지 전처리 ---
+        # 입력 이미지 전처리
         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(LORA_DIR, weight_name=lora_file, adapter_name=adapter_name)
         pipe.set_adapters([adapter_name], [1.0])
 
-        # --- 프롬프트 빌드 ---
-        human_readable_style = style_name.replace("_", " ")
-        prompt = f"Turn this image into the {human_readable_style} style."
+        # 프롬프트 구성
+        readable_style = style_name.replace("_", " ")
+        prompt = f"Turn this image into the {readable_style} style."
         if prompt_suffix and prompt_suffix.strip():
             prompt += f" {prompt_suffix.strip()}"
 
-        gr.Info("Generating styled image… (20‑60 s)")
+        gr.Info("Generating styled image… (20‑60 s)")
 
         result = pipe(
             image=img,
@@ -199,7 +183,7 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
             width=1024,
         )
 
-        # --- LoRA 언로드 & GPU 메모리 해제 ---
+        # LoRA 언로드 & 메모리 정리
         pipe.unload_lora_weights(adapter_name=adapter_name)
         torch.cuda.empty_cache()
 
@@ -211,259 +195,23 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
         return None
 
 # ------------------------------------------------------------------
-# Gradio UI 정의
+# Gradio UI
 # ------------------------------------------------------------------
-with gr.Blocks(title="FLUX.1 Context Style Transfer", theme=gr.themes.Soft()) as demo:
-    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",
-            )
-            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(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=
-"""
-FLUX.1 Kontext Style Transfer
-==============================
-Updated: 2025‑07‑12 (HF_TOKEN 지원)
----------------------------------
-이 스크립트는 Hugging Face **FLUX.1‑Kontext‑dev** 모델과
-22 종의 스타일 LoRA 가중치를 이용해 이미지를 다양한 예술
-스타일로 변환하는 Gradio 데모입니다.
-
-주요 개선 사항
---------------
-1. **HF_TOKEN 환경변수 지원** – `LocalTokenNotFoundError`를
-   방지하기 위해 `HF_TOKEN`(또는 런타임 로그인) 값을 자동으로
-   감지해 `snapshot_download()`에 전달합니다.
-2. **모델 캐싱** – `snapshot_download()`로 실행 시작 시 한 번만
-   모델과 LoRA를 캐싱.
-3. **GPU VRAM 자동 판별** – 24 GB 미만이면 FP16 / CPU offload.
-4. **단일 로딩 메시지** – Gradio `gr.Info()` 메시지가 최초 1회만
-   표시되도록 유지.
-5. **버그 픽스** – seed 처리, LoRA 언로드, 이미지 리사이즈.
-
-------------------------------------------------------------
-"""
-import os
-import gradio as gr
-import spaces
-import torch
-from huggingface_hub import snapshot_download, login as hf_login
-from huggingface_hub.errors import LocalTokenNotFoundError
-from diffusers import FluxKontextPipeline
-from diffusers.utils import load_image
-from PIL import Image
-
-# ------------------------------------------------------------------
-# 환경 설정 & 모델 / 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"))
-
-# --- HF 토큰 처리 --------------------------------------------------
-HF_TOKEN = os.getenv("HF_TOKEN")  # 런타임에 환경변수로 주입하거나
-# docker / Space 설정의 Secrets → HF_TOKEN 로 등록
-
-def _download_with_token(repo_id: str) -> str:
-    """Download repo snapshot with optional token handling."""
-    try:
-        return snapshot_download(
-            repo_id=repo_id,
-            cache_dir=CACHE_DIR,
-            resume_download=True,
-            token=HF_TOKEN if HF_TOKEN else True,  # True → HF_CACHED_TOKEN(or login)
-        )
-    except LocalTokenNotFoundError:
-        # 미로그인 + 필수 동의 모델이면 에러 메시지 출력 후 종료
-        raise RuntimeError(
-            "Hugging Face 토큰이 필요합니다. 환경변수 HF_TOKEN을 설정하거나\n"
-            "`huggingface-cli login`으로 로그인해 주세요."
-        )
+def update_description(style):
+    return STYLE_DESCRIPTIONS.get(style, "")
 
-# --- 최초 실행 시에만 다운로드(이미 캐시에 있으면 건너뜀) ---
-MODEL_DIR = _download_with_token(MODEL_ID)
-LORA_DIR = _download_with_token(LORA_REPO)
 
-# ------------------------------------------------------------------
-# 스타일 → LoRA 파일 매핑 & 설명
-# ------------------------------------------------------------------
-STYLE_LORA_MAP = {
-    "3D_Chibi": "3D_Chibi_lora_weights.safetensors",
-    "American_Cartoon": "American_Cartoon_lora_weights.safetensors",
-    "Chinese_Ink": "Chinese_Ink_lora_weights.safetensors",
-    "Clay_Toy": "Clay_Toy_lora_weights.safetensors",
-    "Fabric": "Fabric_lora_weights.safetensors",
-    "Ghibli": "Ghibli_lora_weights.safetensors",
-    "Irasutoya": "Irasutoya_lora_weights.safetensors",
-    "Jojo": "Jojo_lora_weights.safetensors",
-    "Oil_Painting": "Oil_Painting_lora_weights.safetensors",
-    "Pixel": "Pixel_lora_weights.safetensors",
-    "Snoopy": "Snoopy_lora_weights.safetensors",
-    "Poly": "Poly_lora_weights.safetensors",
-    "LEGO": "LEGO_lora_weights.safetensors",
-    "Origami": "Origami_lora_weights.safetensors",
-    "Pop_Art": "Pop_Art_lora_weights.safetensors",
-    "Van_Gogh": "Van_Gogh_lora_weights.safetensors",
-    "Paper_Cutting": "Paper_Cutting_lora_weights.safetensors",
-    "Line": "Line_lora_weights.safetensors",
-    "Vector": "Vector_lora_weights.safetensors",
-    "Picasso": "Picasso_lora_weights.safetensors",
-    "Macaron": "Macaron_lora_weights.safetensors",
-    "Rick_Morty": "Rick_Morty_lora_weights.safetensors",
-}
+with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🎨 FLUX.1 Kontext Style Transfer
 
-STYLE_DESCRIPTIONS = {
-    "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",
-}
-
-# ------------------------------------------------------------------
-# 파이프라인 로더 (단일 인스턴스)
-# ------------------------------------------------------------------
-_pipeline = None  # 내부 글로벌 캐시
-
-def load_pipeline():
-    """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,
+        업로드한 이미지를 22 종 예술 스타일로 변환하세요!
+        (모델 / LoRA는 최초 실행 시에만 다운로드되며, 이후 실행은 빠릅니다.)
+        """
     )
 
-    pipe.to("cuda")
-
-    if vram_gb < 24:
-        pipe.enable_sequential_cpu_offload()
-    else:
-        pipe.enable_model_cpu_offload()
-
-    _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("Please upload an image first!")
-        return None
-
-    try:
-        pipe = load_pipeline()
-
-        # --- Torch Generator 설정 ---
-        generator = None
-        if seed and int(seed) > 0:
-            generator = torch.Generator(device="cuda").manual_seed(int(seed))
-
-        # --- 입력 이미지 전처리 ---
-        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_file = STYLE_LORA_MAP[style_name]
-        adapter_name = "style"
-        pipe.load_lora_weights(LORA_DIR, weight_name=lora_file, adapter_name=adapter_name)
-        pipe.set_adapters([adapter_name], [1.0])
-
-        # --- 프롬프트 빌드 ---
-        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("Generating styled image… (20‑60 s)")
-
-        result = pipe(
-            image=img,
-            prompt=prompt,
-            guidance_scale=float(guidance_scale),
-            num_inference_steps=int(num_inference_steps),
-            generator=generator,
-            height=1024,
-            width=1024,
-        )
-
-        # --- LoRA 언로드 & GPU 메모리 해제 ---
-        pipe.unload_lora_weights(adapter_name=adapter_name)
-        torch.cuda.empty_cache()
-
-        return result.images[0]
-
-    except Exception as e:
-        torch.cuda.empty_cache()
-        gr.Error(f"Error during style transfer: {e}")
-        return None
-
-# ------------------------------------------------------------------
-# Gradio UI 정의
-# ------------------------------------------------------------------
-with gr.Blocks(title="FLUX.1 Context Style Transfer", theme=gr.themes.Soft()) as demo:
-    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)
@@ -472,15 +220,62 @@ with gr.Blocks(title="FLUX.1 Context Style Transfer", theme=gr.themes.Soft()) as
                 value="Ghibli",
                 label="Select Style",
             )
-            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)
+            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(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")
+                num_steps = gr.Slider(10, 50, value=24, step=1, label="Inference Steps")
+                guidance = gr.Slider(1.0, 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=
+            output_image = gr.Image(label="Styled Result", type="pil", height=400)
+            gr.Markdown(
+                """
+                ### 💡 Tips
+                - 모델(7 GB)·LoRA는 최초 실행 시에만 다운로드됩니다.
+                - 이미지는 1024×1024로 리사이즈 후 처리됩니다.
+                - VRAM < 24 GB인 경우 자동으로 FP16 + CPU offload가 적용됩니다.
+                - seed 값을 변경해 다양한 결과를 얻어 보세요!
+                """
+            )
+
+    # 스타일 설명 자동 업데이트
+    style_dropdown.change(update_description, 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"],
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "Pixel", "8-bit retro game style"],
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "Chinese_Ink", "mountain landscape"],
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "LEGO", "colorful blocks"],
+        ],
+        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,
+    )
+
+    # 버튼 클릭 연결
+    generate_btn.click(
+        fn=style_transfer,
+        inputs=[input_image, style_dropdown, prompt_suffix, num_steps, guidance, seed],
+        outputs=output_image,
+    )
+
+if __name__ == "__main__":
+    demo.launch()