import gradio as gr import spaces import torch from diffusers import FluxKontextPipeline from diffusers.utils import load_image from PIL import Image import os # ---------------------------------------------- # Style → LoRA file mapping # ---------------------------------------------- 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" } # ---------------------------------------------- # 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": "릭 앤 모티 애니메이션 스타일" } # 전역 파이프라인 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 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() return pipe @spaces.GPU(duration=600) # 초기 다운로드 시간 확보 def style_transfer( input_image, style_name, prompt_suffix, num_inference_steps, guidance_scale, seed, ): """선택한 스타일로 이미지 변환""" if input_image is None: gr.Warning("🖼️ 먼저 이미지를 업로드하세요!") 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)) # 이미지 로드 & 리사이즈 img = load_image(input_image) if isinstance(input_image, str) else 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( "Owen777/Kontext-Style-Loras", weight_name=lora_file, adapter_name=adapter_name, ) pipe.set_adapters([adapter_name], adapter_weights=[1.0]) # 프롬프트 구성 prompt = f"Turn this image into the {style_name.replace('_', ' ')} style." if prompt_suffix and prompt_suffix.strip(): prompt += f" {prompt_suffix.strip()}" gr.Info("🎨 이미지를 생성 중…") 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 언로드 및 캐시 정리 pipe.unload_lora_weights() torch.cuda.empty_cache() return result.images[0] except Exception as e: gr.Error(f"🚨 오류: {e}") torch.cuda.empty_cache() return None def update_description(style): return STYLE_DESCRIPTIONS.get(style, "") 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로 이미지를 다양한 예술 스타일로 변환하세요. """ ) 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, ) 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" ) 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, ], outputs=[output_image], ) 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) """ ) if __name__ == "__main__": demo.launch()