Create app.py

app.py

@@ -0,0 +1,202 @@
+import gradio as gr
+import spaces
+import torch
+from huggingface_hub import hf_hub_download
+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"
+}
+
+# Create LoRAs directory if it doesn't exist
+os.makedirs("./LoRAs", exist_ok=True)
+
+# Download all LoRA weights at startup
+print("Downloading LoRA weights...")
+for style_name, lora_file in style_type_lora_dict.items():
+    if not os.path.exists(f"./LoRAs/{lora_file}"):
+        hf_hub_download(
+            repo_id="Owen777/Kontext-Style-Loras", 
+            filename=lora_file, 
+            local_dir="./LoRAs"
+        )
+print("All LoRA weights downloaded!")
+
+# Initialize pipeline globally (will be loaded to GPU when needed)
+pipeline = None
+
+def load_pipeline():
+    global pipeline
+    if pipeline is None:
+        pipeline = FluxKontextPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Kontext-dev", 
+            torch_dtype=torch.bfloat16
+        )
+    return pipeline
+
+@spaces.GPU(duration=120)  # Request GPU for 120 seconds
+def style_transfer(input_image, style_name, num_inference_steps, guidance_scale, seed):
+    """
+    Apply style transfer to the input image using selected style
+    """
+    # Load pipeline and move to GPU
+    pipe = load_pipeline()
+    pipe = pipe.to('cuda')
+    
+    # Set seed for reproducibility
+    if seed is not None and seed > 0:
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+    else:
+        generator = None
+    
+    # Resize input image to 1024x1024
+    if isinstance(input_image, str):
+        image = load_image(input_image)
+    else:
+        image = input_image
+    
+    image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
+    
+    # Load the selected LoRA
+    lora_path = f"./LoRAs/{style_type_lora_dict[style_name]}"
+    pipe.load_lora_weights(lora_path, adapter_name="style_lora")
+    pipe.set_adapters(["style_lora"], adapter_weights=[1.0])
+    
+    # Generate the styled image
+    prompt = f"Turn this image into the {style_name.replace('_', ' ')} style."
+    
+    result = pipe(
+        image=image,
+        prompt=prompt,
+        height=1024,
+        width=1024,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    )
+    
+    # Clear GPU memory
+    torch.cuda.empty_cache()
+    
+    return result.images[0]
+
+# Create Gradio interface
+with gr.Blocks(title="Flux Kontext Style Transfer") as demo:
+    gr.Markdown("""
+    # 🎨 Flux Kontext Style Transfer
+    
+    Transform your images into various artistic styles using FLUX.1-Kontext and style-specific LoRAs.
+    
+    Upload an image and select a style to apply the transformation!
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="Input Image",
+                type="pil",
+                height=400
+            )
+            
+            style_dropdown = gr.Dropdown(
+                choices=list(style_type_lora_dict.keys()),
+                value="3D_Chibi",
+                label="Select Style",
+                info="Choose the artistic style to apply"
+            )
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                num_steps = gr.Slider(
+                    minimum=10,
+                    maximum=50,
+                    value=24,
+                    step=1,
+                    label="Number of Inference Steps",
+                    info="More steps = better quality but slower"
+                )
+                
+                guidance = gr.Slider(
+                    minimum=1.0,
+                    maximum=10.0,
+                    value=3.5,
+                    step=0.5,
+                    label="Guidance Scale",
+                    info="Higher values = stronger style adherence"
+                )
+                
+                seed = gr.Number(
+                    label="Seed",
+                    value=0,
+                    precision=0,
+                    info="Set to 0 for random, or use specific seed for reproducibility"
+                )
+            
+            generate_btn = gr.Button("🎨 Apply Style Transfer", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(
+                label="Styled Output",
+                type="pil",
+                height=400
+            )
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            ["https://huggingface.co/datasets/black-forest-labs/kontext-bench/resolve/main/test/images/0003.jpg", "3D_Chibi"],
+            ["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", "Van_Gogh"],
+        ],
+        inputs=[input_image, style_dropdown],
+        outputs=output_image,
+        fn=lambda img, style: style_transfer(img, style, 24, 3.5, 0),
+        cache_examples=True
+    )
+    
+    # Connect the generate button
+    generate_btn.click(
+        fn=style_transfer,
+        inputs=[input_image, style_dropdown, num_steps, guidance, seed],
+        outputs=output_image
+    )
+    
+    gr.Markdown("""
+    ## 📝 Notes:
+    - Processing takes about 30-60 seconds depending on the number of steps
+    - All images are resized to 1024x1024 for optimal results
+    - Different styles work better with different types of images
+    - Try adjusting the advanced settings for better results
+    
+    ## 🎨 Available Styles:
+    3D Chibi, American Cartoon, Chinese Ink, Clay Toy, Fabric, Ghibli, Irasutoya, 
+    Jojo, Oil Painting, Pixel, Snoopy, Poly, LEGO, Origami, Pop Art, Van Gogh, 
+    Paper Cutting, Line, Vector, Picasso, Macaron, Rick & Morty
+    """)
+
+if __name__ == "__main__":
+    demo.launch()