Update app.py

app.py

@@ -1,60 +1,13 @@
-"""
-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
-from huggingface_hub.errors import LocalTokenNotFoundError
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 from PIL import Image
+import os
 
-# ------------------------------------------------------------------
-# 환경 설정 & 모델 / 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_TOKEN = os.getenv("HF_TOKEN")  # 런타임에 주입하거나 Secrets 사용
-
-
-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 → 로그인 세션 사용
-        )
-    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)
-
-# ------------------------------------------------------------------
-# 스타일 → LoRA 파일 매핑 & 설명
-# ------------------------------------------------------------------
-STYLE_LORA_MAP = {
+# Style dictionary
+style_type_lora_dict = {
     "3D_Chibi": "3D_Chibi_lora_weights.safetensors",
     "American_Cartoon": "American_Cartoon_lora_weights.safetensors",
     "Chinese_Ink": "Chinese_Ink_lora_weights.safetensors",
@@ -76,10 +29,11 @@ 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
+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",
@@ -101,181 +55,233 @@ STYLE_DESCRIPTIONS = {
     "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",
+    "Rick_Morty": "Rick and Morty cartoon style"
 }
 
-# ------------------------------------------------------------------
-# 파이프라인 로더 (싱글턴)
-# ------------------------------------------------------------------
-_pipeline = None
-
+# Initialize pipeline globally
+pipeline = None
+pipeline_loaded = False
 
 def load_pipeline():
-    """Load or return cached FluxKontextPipeline."""
-    global _pipeline
-    if _pipeline is not None:
-        return _pipeline
-
-    # VRAM 판별 → dtype & offload 설정
-    vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
-    dtype = torch.bfloat16 if vram_gb >= 24 else torch.float16
-
-    gr.Info("FLUX.1‑Kontext 파이프라인 로딩 중… (최초 1회)")
-
-    pipe = FluxKontextPipeline.from_pretrained(
-        MODEL_DIR,
-        torch_dtype=dtype,
-        local_files_only=True,
-    ).to("cuda")
-
-    if vram_gb < 24:
-        pipe.enable_sequential_cpu_offload()
-    else:
-        pipe.enable_model_cpu_offload()
-
-    _pipeline = pipe
-    return _pipeline
-
-# ------------------------------------------------------------------
-# 스타일 변환 함수
-# ------------------------------------------------------------------
-@spaces.GPU(duration=600)
+    global pipeline, pipeline_loaded
+    if pipeline is None:
+        print("Loading FLUX.1-Kontext-dev model...")
+        # HF_TOKEN 자동 감지
+        token = os.getenv("HF_TOKEN", True)
+        
+        pipeline = FluxKontextPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Kontext-dev", 
+            torch_dtype=torch.bfloat16,
+            use_auth_token=token
+        )
+        pipeline_loaded = True
+    return pipeline
 
+@spaces.GPU(duration=120)
 def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps, guidance_scale, seed):
-    """Apply selected style to the uploaded image."""
+    """
+    Apply style transfer to the input image using selected style
+    """
     if input_image is None:
         gr.Warning("Please upload an image first!")
         return None
-
+    
     try:
+        # Load pipeline and move to GPU
         pipe = load_pipeline()
-
-        # Torch Generator (seed 고정 시 재현 가능)
+        pipe = pipe.to('cuda')
+        
+        # Enable memory efficient settings
+        pipe.enable_model_cpu_offload()
+        
+        # Set seed for reproducibility
         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])
-
-        # 프롬프트 구성
-        readable_style = style_name.replace("_", " ")
-        prompt = f"Turn this image into the {readable_style} style."
+        if seed > 0:
+            generator = torch.Generator(device="cuda").manual_seed(seed)
+        
+        # Process input image
+        if isinstance(input_image, str):
+            image = load_image(input_image)
+        else:
+            image = input_image
+        
+        # Ensure RGB and resize to 1024x1024
+        image = image.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS)
+        
+        # Load the selected LoRA
+        lora_filename = style_type_lora_dict[style_name]
+        
+        # Clear any previously loaded LoRA
+        try:
+            pipe.unload_lora_weights()
+        except:
+            pass
+        
+        # Load LoRA weights
+        pipe.load_lora_weights(
+            "Owen777/Kontext-Style-Loras",
+            weight_name=lora_filename,
+            adapter_name="style"
+        )
+        pipe.set_adapters(["style"], adapter_weights=[1.0])
+        
+        # Create prompt for style transformation
+        style_name_readable = style_name.replace('_', ' ')
+        prompt = f"Turn this image into the {style_name_readable} style."
         if prompt_suffix and prompt_suffix.strip():
             prompt += f" {prompt_suffix.strip()}"
-
-        gr.Info("Generating styled image… (20‑60 s)")
-
+        
+        print(f"Generating with prompt: {prompt}")
+        
+        # Generate the styled image
         result = pipe(
-            image=img,
+            image=image,
             prompt=prompt,
-            guidance_scale=float(guidance_scale),
-            num_inference_steps=int(num_inference_steps),
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
             generator=generator,
             height=1024,
-            width=1024,
+            width=1024
         )
-
-        # LoRA 언로드 & 메모리 정리
-        pipe.unload_lora_weights(adapter_name=adapter_name)
+        
+        # Clear GPU memory
         torch.cuda.empty_cache()
-
+        
         return result.images[0]
-
+        
     except Exception as e:
+        print(f"Error: {str(e)}")
+        gr.Error(f"Error during style transfer: {str(e)}")
         torch.cuda.empty_cache()
-        gr.Error(f"Error during style transfer: {e}")
         return None
 
-# ------------------------------------------------------------------
-# Gradio UI
-# ------------------------------------------------------------------
-
-def update_description(style):
-    return STYLE_DESCRIPTIONS.get(style, "")
-
-
+# Create Gradio interface
 with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🎨 FLUX.1 Kontext Style Transfer
-
-        업로드한 이미지를 22 종 예술 스타일로 변환하세요!
-        (모델 / LoRA는 최초 실행 시에만 다운로드되며, 이후 실행은 빠릅니다.)
-        """
-    )
-
+    gr.Markdown("""
+    # 🎨 FLUX.1 Kontext Style Transfer
+    
+    Transform your images into various artistic styles using FLUX.1-Kontext-dev and high-quality style LoRAs.
+    
+    This demo uses the official Owen777/Kontext-Style-Loras collection with 22 different artistic styles!
+    """)
+    
     with gr.Row():
         with gr.Column(scale=1):
-            input_image = gr.Image(label="Upload Image", type="pil", height=400)
+            input_image = gr.Image(
+                label="Upload Image",
+                type="pil",
+                height=400
+            )
+            
             style_dropdown = gr.Dropdown(
-                choices=list(STYLE_LORA_MAP.keys()),
+                choices=list(style_type_lora_dict.keys()),
                 value="Ghibli",
                 label="Select Style",
+                info="Choose from 22 different artistic styles"
             )
+            
             style_info = gr.Textbox(
                 label="Style Description",
-                value=STYLE_DESCRIPTIONS["Ghibli"],
+                value=style_descriptions["Ghibli"],
                 interactive=False,
-                lines=2,
+                lines=2
             )
+            
             prompt_suffix = gr.Textbox(
                 label="Additional Instructions (Optional)",
-                placeholder="e.g. add dramatic lighting",
-                lines=2,
+                placeholder="Add extra details like 'make it more colorful' or '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")
-                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)
-
+                num_steps = gr.Slider(
+                    minimum=10,
+                    maximum=50,
+                    value=24,
+                    step=1,
+                    label="Inference Steps",
+                    info="More steps = better quality but slower"
+                )
+                
+                guidance = gr.Slider(
+                    minimum=1.0,
+                    maximum=5.0,
+                    value=2.5,
+                    step=0.1,
+                    label="Guidance Scale",
+                    info="How closely to follow the prompt (2.5 recommended)"
+                )
+                
+                seed = gr.Number(
+                    label="Seed",
+                    value=42,
+                    precision=0,
+                    info="Set to 0 for random results"
+                )
+            
             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
-                - 모델(7 GB)·LoRA는 최초 실행 시에만 다운로드됩니다.
-                - 이미지는 1024×1024로 리사이즈 후 처리됩니다.
-                - VRAM < 24 GB인 경우 자동으로 FP16 + CPU offload가 적용됩니다.
-                - seed 값을 변경해 다양한 결과를 얻어 보세요!
-                """
+            output_image = gr.Image(
+                label="Styled Result",
+                type="pil",
+                height=400
             )
-
-    # 스타일 설명 자동 업데이트
-    style_dropdown.change(update_description, inputs=[style_dropdown], outputs=[style_info])
-
-    # 예제 샘플
+            
+            gr.Markdown("""
+            ### 💡 Tips:
+            - All images are resized to 1024x1024
+            - First run downloads the model (~12GB)
+            - Each style transformation takes ~30-60 seconds
+            - Try different styles to find the best match!
+            - Use additional instructions for fine control
+            """)
+    
+    # Update style description when style changes
+    def update_description(style):
+        return style_descriptions.get(style, "")
+    
+    style_dropdown.change(
+        fn=update_description,
+        inputs=[style_dropdown],
+        outputs=[style_info]
+    )
+    
+    # Examples
     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,
+        fn=style_transfer,
+        cache_examples=False
     )
-
-    # 버튼 클릭 연결
+    
+    # Connect the generate button
     generate_btn.click(
         fn=style_transfer,
         inputs=[input_image, style_dropdown, prompt_suffix, num_steps, guidance, seed],
-        outputs=output_image,
+        outputs=output_image
     )
+    
+    gr.Markdown("""
+    ---
+    ### 📚 Available Styles:
+    
+    **Anime/Cartoon**: Ghibli, American Cartoon, Jojo, Snoopy, Rick & Morty, Irasutoya  
+    **3D/Geometric**: 3D Chibi, Poly, LEGO, Clay Toy  
+    **Traditional Art**: Chinese Ink, Oil Painting, Van Gogh, Picasso, Pop Art  
+    **Craft/Material**: Fabric, Origami, Paper Cutting, Macaron  
+    **Digital/Modern**: Pixel, Line, Vector
+    
+    ---
+    
+    Created with ❤️ using [Owen777/Kontext-Style-Loras](https://huggingface.co/Owen777/Kontext-Style-Loras)
+    """)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()