Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import spaces
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 from PIL import Image
 import os
@@ -39,15 +39,15 @@ 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(
+        # Load FLUX.1-Kontext-dev model
+        pipeline = FluxKontextPipeline.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):
+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
     """
@@ -60,14 +60,14 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
         pipe = load_pipeline()
         pipe = pipe.to('cuda')
         
-        # Enable memory efficient attention
+        # Enable memory efficient settings
         pipe.enable_model_cpu_offload()
         
         # Set seed for reproducibility
         if seed > 0:
             generator = torch.Generator(device="cuda").manual_seed(seed)
         else:
-            generator = None
+            generator = torch.manual_seed(42)
         
         # Process input image
         if isinstance(input_image, str):
@@ -76,40 +76,39 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
             image = input_image
         
         # Resize to 1024x1024 (required for Kontext)
-        image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
+        image = image.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS)
         
-        # Load the selected LoRA from the repository
+        # Load the selected LoRA
         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"
+            weight_name=lora_filename
         )
-        pipe.set_adapters(["style"], adapter_weights=[1.0])
+        pipe.set_adapters(["default"], adapter_weights=[1.0])
         
-        # Create prompt
+        # 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:
-            prompt += f" {prompt_suffix}"
+        if prompt_suffix and prompt_suffix.strip():
+            prompt += f" {prompt_suffix.strip()}"
         
         gr.Info("Generating styled image...")
         
-        # Generate the styled image
+        # Generate the styled image with Kontext pipeline
         result = pipe(
             image=image,
             prompt=prompt,
-            height=1024,
-            width=1024,
+            guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             generator=generator,
-            guidance_scale=3.5
+            height=1024,
+            width=1024
         )
         
-        # Clear GPU memory
+        # Clear LoRA and GPU memory
         pipe.unload_lora_weights()
         torch.cuda.empty_cache()
         
@@ -179,8 +178,8 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
             )
             
             prompt_suffix = gr.Textbox(
-                label="Additional Prompt (Optional)",
-                placeholder="Add extra details to the transformation...",
+                label="Additional Instructions (Optional)",
+                placeholder="Add extra details like 'make it more colorful' or 'add dramatic lighting'...",
                 lines=2
             )
             
@@ -194,6 +193,15 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
                     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,
@@ -215,6 +223,7 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
             - First run downloads the model (~7GB)
             - 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
@@ -239,14 +248,14 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
         ],
         inputs=[input_image, style_dropdown, prompt_suffix],
         outputs=output_image,
-        fn=lambda img, style, prompt: style_transfer(img, style, prompt, 24, 42),
+        fn=lambda img, style, prompt: style_transfer(img, style, prompt, 24, 2.5, 42),
         cache_examples=False
     )
     
     # Connect the generate button
     generate_btn.click(
         fn=style_transfer,
-        inputs=[input_image, style_dropdown, prompt_suffix, num_steps, seed],
+        inputs=[input_image, style_dropdown, prompt_suffix, num_steps, guidance, seed],
         outputs=output_image
     )
     
@@ -269,7 +278,7 @@ 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
+    3. (Optional) Add custom instructions
     4. Click "Transform Image" and wait ~30-60 seconds
     5. Download your styled image!