app.py

import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline
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"
}

# Initialize pipeline globally
pipeline = None

def load_pipeline():
    global pipeline
    if pipeline is None:
        gr.Info("Loading FLUX.1-Kontext model...")
        # Use DiffusionPipeline to load FLUX.1-Kontext-dev
        pipeline = DiffusionPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-Kontext-dev", 
            torch_dtype=torch.bfloat16
        )
    return pipeline

@spaces.GPU(duration=120)
def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps, 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 attention
        pipe.enable_model_cpu_offload()
        
        # Set seed for reproducibility
        if seed > 0:
            generator = torch.Generator(device="cuda").manual_seed(seed)
        else:
            generator = None
        
        # Process input image
        if isinstance(input_image, str):
            image = load_image(input_image)
        else:
            image = input_image
        
        # Resize to 1024x1024 (required for Kontext)
        image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
        
        # Load the selected LoRA from the repository
        gr.Info(f"Loading {style_name} style...")
        lora_filename = style_type_lora_dict[style_name]
        
        # Load LoRA weights directly from the repository
        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
        style_name_readable = style_name.replace('_', ' ')
        prompt = f"Turn this image into the {style_name_readable} style."
        if prompt_suffix:
            prompt += f" {prompt_suffix}"
        
        gr.Info("Generating styled image...")
        
        # Generate the styled image
        result = pipe(
            image=image,
            prompt=prompt,
            height=1024,
            width=1024,
            num_inference_steps=num_inference_steps,
            generator=generator,
            guidance_scale=3.5
        )
        
        # Clear GPU memory
        pipe.unload_lora_weights()
        torch.cuda.empty_cache()
        
        return result.images[0]
        
    except Exception as e:
        gr.Error(f"Error during style transfer: {str(e)}")
        torch.cuda.empty_cache()
        return None

# 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"
}

# 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!
    """)
    
    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="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 Prompt (Optional)",
                placeholder="Add extra details to the transformation...",
                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"
                )
                
                seed = gr.Number(
                    label="Seed",
                    value=42,
                    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:
            - All images are resized to 1024x1024
            - First run downloads the model (~7GB)
            - Each style transformation takes ~30-60 seconds
            - Try different styles to find the best match!
            """)
    
    # 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, 42),
        cache_examples=False
    )
    
    # Connect the generate button
    generate_btn.click(
        fn=style_transfer,
        inputs=[input_image, style_dropdown, prompt_suffix, num_steps, 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
    
    ---
    
    ### 🚀 How it works:
    1. Upload any image
    2. Select a style from the dropdown
    3. (Optional) Add custom prompt details
    4. Click "Transform Image" and wait ~30-60 seconds
    5. Download your styled image!
    
    ---
    
    Created with ❤️ using [Owen777/Kontext-Style-Loras](https://huggingface.co/Owen777/Kontext-Style-Loras)
    """)

if __name__ == "__main__":
    demo.launch()