Update app.py

app.py

@@ -20,7 +20,7 @@ pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev
 
 # Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)
 
-with open("flux_loras.json", "r") as file:
+with open("flux_loras-v2.json", "r") as file:
     data = json.load(file)
     flux_loras_raw = [
         {
@@ -30,6 +30,7 @@ with open("flux_loras.json", "r") as file:
             "trigger_word": item.get("trigger_word", ""),
             "trigger_position": item.get("trigger_position", "prepend"),
             "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
+            "lora_type": item.get("lora_type", "flux"),
         }
         for item in data
     ]
@@ -126,12 +127,12 @@ def classify_gallery(flux_loras):
     sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
     return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
 
-def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75,portrait_mode=False, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
     """Wrapper function to handle state serialization"""
     return infer_with_lora(input_image, prompt, selected_index, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, flux_loras, progress)
 
 @spaces.GPU
-def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, portrait_mode=False, flux_loras=None,  progress=gr.Progress(track_tqdm=True)):
     """Generate image with selected LoRA"""
     global current_lora, pipe
     
@@ -169,13 +170,19 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
     input_image = input_image.convert("RGB")
     # Add trigger word to prompt
     trigger_word = lora_to_use["trigger_word"]
-    if trigger_word == ", How2Draw":
-        prompt = f"create a How2Draw sketch of the person of the photo {prompt}, maintain the facial identity of the person and general features"
-    elif trigger_word == ", video game screenshot in the style of THSMS":
-        prompt = f"create a video game screenshot in the style of THSMS with the person from the photo, {prompt}. maintain the facial identity of the person and general features"
-    else:
-        prompt = f"convert the style of this portrait photo to {trigger_word} while maintaining the identity of the person. {prompt}. Make sure to maintain the person's facial identity and features, while still changing the overall style to {trigger_word}."
-    
+    is_kontext_lora = lora_to_use["lora_type"] == "kontext"
+    if not is_kontext_lora:
+        if trigger_word == ", How2Draw":
+            prompt = f"create a How2Draw sketch of the person of the photo {prompt}, maintain the facial identity of the person and general features"
+        elif trigger_word == ", video game screenshot in the style of THSMS":
+            prompt = f"create a video game screenshot in the style of THSMS with the person from the photo, {prompt}. maintain the facial identity of the person and general features"
+        else:
+            prompt = f"convert the style of this portrait photo to {trigger_word} while maintaining the identity of the person. {prompt}. Make sure to maintain the person's facial identity and features, while still changing the overall style to {trigger_word}."
+    else: 
+        if portrait_mode:
+            prompt = f"{trigger_word} while maintaining the identity of the person. {prompt}. Make sure to maintain the person's facial identity and features."
+        else:
+            prompt = f"{trigger_word}. {prompt}."
     try:
         image = pipe(
             image=input_image, 
@@ -242,7 +249,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(scale=4, elem_id="box_column"):
             with gr.Group(elem_id="gallery_box"):
                 input_image = gr.Image(label="Upload a picture of yourself", type="pil", height=300)
-                
+                portrait_mode = gr.Checkbox(label="portrait mode", value=True)
                 gallery = gr.Gallery(
                     label="Pick a LoRA",
                     allow_preview=False,
@@ -328,7 +335,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer_with_lora_wrapper,
-        inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, gr_flux_loras],
+        inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, gr_flux_loras],
         outputs=[result, seed, reuse_button]
     )