Create app.py

app.py

@@ -0,0 +1,357 @@
+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 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",
+    "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": "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"
+}
+
+# Mapping for thumbnail files
+thumbnail_mapping = {
+    "3D_Chibi": "3D_Chibi.webp",
+    "American_Cartoon": "american_cartoon.webp",
+    "Chinese_Ink": "chinese_ink.webp",
+    "Clay_Toy": "clay_toy.webp",
+    "Fabric": "fabric.webp",
+    "Ghibli": "ghibli.webp",
+    "Irasutoya": "Irasutoya.webp",
+    "Jojo": "jojo.webp",
+    "Oil_Painting": "oil_painting.webp",
+    "Pixel": "pixel.webp",
+    "Snoopy": "snoopy.webp",
+    "Poly": "poly.webp",
+    "LEGO": "LEGO.webp",
+    "Origami": "origami.webp",
+    "Pop_Art": "pop-art.webp",
+    "Van_Gogh": "van_gogh.webp",
+    "Paper_Cutting": "Paper_Cutting.webp",
+    "Line": "line.webp",
+    "Vector": "vector.webp",
+    "Picasso": "picasso.webp",
+    "Macaron": "Macaron.webp",
+    "Rick_Morty": "Rick_Morty.webp"
+}
+
+# Initialize pipeline globally
+pipeline = None
+pipeline_loaded = False
+
+def load_pipeline():
+    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 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()
+        pipe = pipe.to('cuda')
+        
+        # Enable memory efficient settings
+        pipe.enable_model_cpu_offload()
+        
+        # Set seed for reproducibility
+        generator = None
+        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()}"
+        
+        print(f"Generating with prompt: {prompt}")
+        
+        # Generate the styled image
+        result = pipe(
+            image=image,
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            height=1024,
+            width=1024
+        )
+        
+        # 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()
+        return None
+
+def create_thumbnail_html():
+    """Create HTML for thumbnail grid"""
+    html = '<div style="display: grid; grid-template-columns: repeat(6, 1fr); gap: 10px; max-width: 800px; margin: 0 auto;">'
+    
+    styles = list(style_type_lora_dict.keys())
+    
+    for i, style in enumerate(styles):
+        if i >= 24:  # Limit to 24 thumbnails for 6x4 grid
+            break
+        
+        thumbnail_file = thumbnail_mapping.get(style, "")
+        style_readable = style.replace('_', ' ')
+        
+        html += f'''
+        <div style="text-align: center; cursor: pointer;" onclick="document.getElementById('style_dropdown').value='{style}'; 
+                     var event = new Event('change', {{bubbles: true}});
+                     document.getElementById('style_dropdown').dispatchEvent(event);">
+            <img src="file/{thumbnail_file}" alt="{style_readable}" 
+                 style="width: 100%; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); 
+                        transition: transform 0.2s, box-shadow 0.2s;"
+                 onmouseover="this.style.transform='scale(1.05)'; this.style.boxShadow='0 4px 8px rgba(0,0,0,0.2)';"
+                 onmouseout="this.style.transform='scale(1)'; this.style.boxShadow='0 2px 4px rgba(0,0,0,0.1)';">
+            <p style="margin: 5px 0; font-size: 12px; font-weight: 500;">{style_readable}</p>
+        </div>
+        '''
+    
+    # Fill empty slots if needed
+    remaining_slots = 24 - len(styles)
+    if remaining_slots > 0 and len(styles) < 24:
+        for _ in range(remaining_slots):
+            html += '<div></div>'
+    
+    html += '</div>'
+    return html
+
+# 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
+    
+    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!
+    """)
+    
+    # Thumbnail Grid Section
+    gr.Markdown("### 🖼️ Click a style thumbnail to select it:")
+    with gr.Row():
+        gr.HTML(create_thumbnail_html())
+    
+    gr.Markdown("---")
+    
+    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_type_lora_dict.keys()),
+                value="Ghibli",
+                label="Selected Style",
+                info="Choose from 22 different artistic styles or click a thumbnail above",
+                elem_id="style_dropdown"
+            )
+            
+            style_info = gr.Textbox(
+                label="Style Description",
+                value=style_descriptions["Ghibli"],
+                interactive=False,
+                lines=2
+            )
+            
+            prompt_suffix = gr.Textbox(
+                label="Additional Instructions (Optional)",
+                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(
+                    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:
+            - Click any thumbnail above to quickly select a style
+            - 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"],
+        ],
+        inputs=[input_image, style_dropdown, prompt_suffix],
+        outputs=output_image,
+        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
+    )
+    
+    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()