Update app.py

app.py

@@ -1,14 +1,12 @@
 import gradio as gr
 import spaces
 import torch
-from huggingface_hub import hf_hub_download
-from diffusers import FluxKontextPipeline
+from diffusers import DiffusionPipeline
 from diffusers.utils import load_image
 from PIL import Image
 import os
-import numpy as np
 
-# Style dictionary - 완전한 스타일 목록
+# Style dictionary
 style_type_lora_dict = {
     "3D_Chibi": "3D_Chibi_lora_weights.safetensors",
     "American_Cartoon": "American_Cartoon_lora_weights.safetensors",
@@ -34,29 +32,6 @@ style_type_lora_dict = {
     "Rick_Morty": "Rick_Morty_lora_weights.safetensors"
 }
 
-# Create LoRAs directory if it doesn't exist
-os.makedirs("./LoRAs", exist_ok=True)
-
-# Download LoRA weights on demand
-def download_lora(style_name):
-    lora_file = style_type_lora_dict[style_name]
-    lora_path = f"./LoRAs/{lora_file}"
-    
-    if not os.path.exists(lora_path):
-        gr.Info(f"Downloading {style_name} LoRA...")
-        try:
-            hf_hub_download(
-                repo_id="Owen777/Kontext-Style-Loras", 
-                filename=lora_file, 
-                local_dir="./LoRAs"
-            )
-            print(f"Downloaded {lora_file}")
-        except Exception as e:
-            print(f"Error downloading {lora_file}: {e}")
-            raise e
-    
-    return lora_path
-
 # Initialize pipeline globally
 pipeline = None
 
@@ -64,13 +39,14 @@ def load_pipeline():
     global pipeline
     if pipeline is None:
         gr.Info("Loading FLUX.1-Kontext model...")
-        pipeline = FluxKontextPipeline.from_pretrained(
+        # 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)  # Request GPU for 120 seconds
+@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
@@ -84,6 +60,9 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
         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)
@@ -99,12 +78,17 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
         # Resize to 1024x1024 (required for Kontext)
         image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
         
-        # Download and load the selected LoRA
+        # Load the selected LoRA from the repository
         gr.Info(f"Loading {style_name} style...")
-        lora_path = download_lora(style_name)
+        lora_filename = style_type_lora_dict[style_name]
         
-        pipe.load_lora_weights(lora_path, adapter_name="style")
-        pipe.set_adapters(["style"], adapter_weights=[1])
+        # 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('_', ' ')
@@ -121,10 +105,12 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
             height=1024,
             width=1024,
             num_inference_steps=num_inference_steps,
-            generator=generator
+            generator=generator,
+            guidance_scale=3.5
         )
         
         # Clear GPU memory
+        pipe.unload_lora_weights()
         torch.cuda.empty_cache()
         
         return result.images[0]
@@ -226,8 +212,8 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
             gr.Markdown("""
             ### 💡 Tips:
             - All images are resized to 1024x1024
-            - First run may take longer to download the model
-            - Each style LoRA is ~359MB and downloaded on first use
+            - First run downloads the model (~7GB)
+            - Each style transformation takes ~30-60 seconds
             - Try different styles to find the best match!
             """)
     
@@ -248,6 +234,8 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
             ["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,
@@ -278,6 +266,15 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
     
     ---
     
+    ### 🚀 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)
     """)