Update app.py

app.py

@@ -18,69 +18,66 @@ MAX_SEED = np.iinfo(np.int32).max
 
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
-# Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)
-try:
-    with open("flux_loras.json", "r") as file:
-        data = json.load(file)
-        flux_loras_raw = [
-            {
-                "image": item["image"],
-                "title": item["title"],
-                "repo": item["repo"],
-                "trigger_word": item.get("trigger_word", ""),
-                "trigger_position": item.get("trigger_position", "prepend"),
-                "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
-                "likes": item.get("likes", 0),
-            }
-            for item in data
-        ]
-    print(f"Successfully loaded {len(flux_loras_raw)} LoRAs from JSON")
-except Exception as e:
-    print(f"Error loading flux_loras.json: {e}")
-    print("Using sample LoRA data instead...")
-    # Sample LoRA data with working repositories
-    flux_loras_raw = [
-        {
-            "image": "https://huggingface.co/alvdansen/flux-koda/resolve/main/images/photo-1586902197503-e71026292412.jpeg",
-            "title": "Flux Koda",
-            "repo": "alvdansen/flux-koda",
-            "trigger_word": "flmft style",
-            "weights": "flux_lora.safetensors",
-            "likes": 100
-        },
-        {
-            "image": "https://huggingface.co/multimodalart/flux-tarot-v1/resolve/main/images/e5f2761e5a474e52ab11b1c9246c9a30.png",
-            "title": "Tarot Cards",
-            "repo": "multimodalart/flux-tarot-v1",
-            "trigger_word": "in the style of TOK a trtcrd tarot style",
-            "weights": "flux_tarot_v1_lora.safetensors",
-            "likes": 90
-        },
-        {
-            "image": "https://huggingface.co/Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style/resolve/main/d13591878de740648a8f29b836e16ff2.jpeg",
-            "title": "Paper Cutout",
-            "repo": "Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style",
-            "trigger_word": "Paper Cutout Style",
-            "weights": "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors",
-            "likes": 80
-        },
-        {
-            "image": "https://huggingface.co/alvdansen/frosting_lane_flux/resolve/main/images/content%20-%202024-08-11T010011.238.jpeg",
-            "title": "Frosting Lane",
-            "repo": "alvdansen/frosting_lane_flux",
-            "trigger_word": "frstingln illustration",
-            "weights": "flux_lora_frosting_lane_flux_000002500.safetensors",
-            "likes": 70
-        },
-        {
-            "image": "https://huggingface.co/davisbro/flux-watercolor/resolve/main/images/wc2.png",
-            "title": "Watercolor",
-            "repo": "davisbro/flux-watercolor",
-            "trigger_word": "watercolor style",
-            "weights": "flux_watercolor.safetensors",
-            "likes": 60
-        }
-    ]
+# Load LoRA data
+flux_loras_raw = [
+    {
+        "image": "examples/1.png",
+        "title": "Studio Ghibli",
+        "repo": "openfree/flux-chatgpt-ghibli-lora",
+        "trigger_word": "ghibli",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    },
+    {
+        "image": "examples/2.png",
+        "title": "Winslow Homer",
+        "repo": "openfree/winslow-homer",
+        "trigger_word": "homer",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    },
+    {
+        "image": "examples/3.png",
+        "title": "Van Gogh",
+        "repo": "openfree/van-gogh",
+        "trigger_word": "gogh",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    },
+    {
+        "image": "examples/4.png",
+        "title": "Paul Cézanne",
+        "repo": "openfree/paul-cezanne",
+        "trigger_word": "Cezanne",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    },
+    {
+        "image": "examples/5.png",
+        "title": "Renoir",
+        "repo": "openfree/pierre-auguste-renoir",
+        "trigger_word": "Renoir",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    },
+    {
+        "image": "examples/6.png",
+        "title": "Claude Monet",
+        "repo": "openfree/claude-monet",
+        "trigger_word": "claude monet",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    },
+    {
+        "image": "examples/7.png",
+        "title": "Fantasy Art",
+        "repo": "openfree/myt-flux-fantasy",
+        "trigger_word": "fantasy",
+        "weights": "pytorch_lora_weights.safetensors",
+        "likes": 0
+    }
+]
+print(f"Loaded {len(flux_loras_raw)} LoRAs")
 # Global variables for LoRA management
 current_lora = None
 lora_cache = {}
@@ -136,11 +133,14 @@ def update_selection(selected_state: gr.SelectData, flux_loras):
     if selected_state.index >= len(flux_loras):
         return "### No LoRA selected", gr.update(), None
     
-    lora_repo = flux_loras[selected_state.index]["repo"]
-    trigger_word = flux_loras[selected_state.index]["trigger_word"]
+    lora = flux_loras[selected_state.index]
+    lora_title = lora["title"]
+    lora_repo = lora["repo"]
+    trigger_word = lora["trigger_word"]
     
-    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
-    new_placeholder = f"optional description, e.g. 'a man with glasses and a beard'"
+    # Create a more informative selected text
+    updated_text = f"### 🎨 Selected Style: {lora_title}"
+    new_placeholder = f"Describe additional details, e.g., 'wearing a red hat' or 'smiling'"
     
     return updated_text, gr.update(placeholder=new_placeholder), selected_state.index
 
@@ -188,17 +188,20 @@ def get_huggingface_lora(link):
 def load_custom_lora(link):
     """Load custom LoRA from user input"""
     if not link:
-        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### Click on a LoRA in the gallery to select it", None
+        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### 🎨 Select an art style from the gallery", None
     
     try:
         repo_name, weights_file, trigger_word = get_huggingface_lora(link)
         
         card = f'''
-        <div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px; margin: 10px 0;">
-            <span><strong>Loaded custom LoRA:</strong></span>
-            <div style="margin-top: 8px;">
-                <h4>{repo_name}</h4>
-                <small>{"Using: <code><b>"+trigger_word+"</b></code> as trigger word" if trigger_word else "No trigger word found"}</small>
+        <div class="custom_lora_card">
+            <div style="display: flex; align-items: center; margin-bottom: 12px;">
+                <span style="font-size: 18px; margin-right: 8px;">✅</span>
+                <strong style="font-size: 16px;">Custom LoRA Loaded!</strong>
+            </div>
+            <div style="background: rgba(255, 255, 255, 0.8); padding: 12px; border-radius: 8px;">
+                <h4 style="margin: 0 0 8px 0; color: #333;">{repo_name}</h4>
+                <small style="color: #666;">{"Trigger: <code style='background: #f0f0f0; padding: 2px 6px; border-radius: 4px;'><b>"+trigger_word+"</b></code>" if trigger_word else "No trigger word found"}</small>
             </div>
         </div>
         '''
@@ -209,10 +212,10 @@ def load_custom_lora(link):
             "trigger_word": trigger_word
         }
         
-        return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"Custom: {repo_name}", None
+        return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"🎨 Custom Style: {repo_name}", None
     
     except Exception as e:
-        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### Click on a LoRA in the gallery to select it", None
+        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### 🎨 Select an art style from the gallery", None
 
 def remove_custom_lora():
     """Remove custom LoRA"""
@@ -224,16 +227,26 @@ def classify_gallery(flux_loras):
         sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
         gallery_items = []
         
+        # Color mapping for each style
+        style_colors = {
+            "Studio Ghibli": "81C784/FFFFFF",  # Green
+            "Winslow Homer": "64B5F6/FFFFFF",  # Blue
+            "Van Gogh": "FFD54F/333333",       # Yellow
+            "Paul Cézanne": "FF8A65/FFFFFF",   # Orange
+            "Renoir": "F06292/FFFFFF",         # Pink
+            "Claude Monet": "9575CD/FFFFFF",   # Purple
+            "Fantasy Art": "4FC3F7/FFFFFF"     # Light Blue
+        }
+        
         for item in sorted_gallery:
             if "image" in item and "title" in item:
                 image_url = item["image"]
                 title = item["title"]
                 
-                # If image is a local file path that might not exist, use a placeholder URL
-                if isinstance(image_url, str) and (image_url.startswith("/home/") or image_url.startswith("samples/") or not image_url.startswith("http")):
-                    print(f"Replacing local/invalid image path: {image_url}")
-                    # Use a more reliable placeholder
-                    image_url = f"https://via.placeholder.com/512x512/E0E7FF/818CF8?text={title.replace(' ', '+')}"
+                # If image is a local file path, use a styled placeholder
+                if isinstance(image_url, str) and (image_url.startswith("examples/") or image_url.startswith("/home/") or image_url.startswith("samples/") or not image_url.startswith("http")):
+                    color = style_colors.get(title, "E0E7FF/818CF8")
+                    image_url = f"https://via.placeholder.com/512x512/{color}?text={title.replace(' ', '+')}"
                     
                 gallery_items.append((image_url, title))
         
@@ -246,7 +259,7 @@ def classify_gallery(flux_loras):
         print(f"Error in classify_gallery: {e}")
         return [], []
 
-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.0, 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)
 
@@ -257,7 +270,7 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
     
     # Check if input image is provided
     if input_image is None:
-        gr.Warning("Please upload an image first!")
+        gr.Warning("Please upload your portrait photo first! 📸")
         return None, seed, gr.update(visible=False)
     
     if randomize_seed:
@@ -290,7 +303,7 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
                 current_lora = lora_to_use
             else:
                 print(f"Failed to load LoRA from {repo_id}")
-                gr.Warning(f"Failed to load LoRA style. Please try a different one.")
+                gr.Warning(f"Failed to load {lora_to_use.get('title', 'style')}. Please try a different art style.")
                 return None, seed, gr.update(visible=False)
             
         except Exception as e:
@@ -305,17 +318,33 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
         input_image = input_image.convert("RGB")
     except Exception as e:
         print(f"Error processing image: {e}")
-        gr.Warning("Error processing the uploaded image. Please try a different image.")
+        gr.Warning("Error processing the uploaded image. Please try a different photo. 📸")
         return None, seed, gr.update(visible=False)
     
     # Check if LoRA is selected
     if lora_to_use is None:
-        gr.Warning("Please select a LoRA style from the gallery first!")
+        gr.Warning("Please select an art style from the gallery first! 🎨")
         return None, seed, gr.update(visible=False)
     
     # Add trigger word to prompt
     trigger_word = lora_to_use.get("trigger_word", "")
-    if trigger_word == ", How2Draw":
+    
+    # Special handling for different art styles
+    if trigger_word == "ghibli":
+        prompt = f"Create a Studio Ghibli anime style portrait of the person in the photo, {prompt}. Maintain the facial identity while transforming into whimsical anime art style."
+    elif trigger_word == "homer":
+        prompt = f"Paint the person in Winslow Homer's American realist style, {prompt}. Keep facial features while applying watercolor and marine art techniques."
+    elif trigger_word == "gogh":
+        prompt = f"Transform the portrait into Van Gogh's post-impressionist style with swirling brushstrokes, {prompt}. Maintain facial identity with expressive colors."
+    elif trigger_word == "Cezanne":
+        prompt = f"Render the person in Paul Cézanne's geometric post-impressionist style, {prompt}. Keep facial structure while applying structured brushwork."
+    elif trigger_word == "Renoir":
+        prompt = f"Paint the portrait in Pierre-Auguste Renoir's impressionist style with soft light, {prompt}. Maintain identity with luminous skin tones."
+    elif trigger_word == "claude monet":
+        prompt = f"Create an impressionist portrait in Claude Monet's style with visible brushstrokes, {prompt}. Keep facial features while using light and color."
+    elif trigger_word == "fantasy":
+        prompt = f"Transform into an epic fantasy character portrait, {prompt}. Maintain facial identity while adding magical and fantastical elements."
+    elif 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"
@@ -571,10 +600,10 @@ with gr.Blocks(css=css) as demo:
     with gr.Row(elem_id="main_app"):
         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)
+                input_image = gr.Image(label="Upload your portrait photo 📸", type="pil", height=300)
                 
                 gallery = gr.Gallery(
-                    label="Pick a LoRA",
+                    label="Choose Your Art Style",
                     allow_preview=False,
                     columns=3,
                     elem_id="gallery",
@@ -583,55 +612,57 @@ with gr.Blocks(css=css) as demo:
                 )
                 
                 custom_model = gr.Textbox(
-                    label="Or enter a custom HuggingFace FLUX LoRA", 
+                    label="🔗 Or use a custom LoRA from HuggingFace", 
                     placeholder="e.g., username/lora-name",
                     visible=True
                 )
                 custom_model_card = gr.HTML(visible=False)
-                custom_model_button = gr.Button("Remove custom LoRA", visible=False)
+                custom_model_button = gr.Button("❌ Remove custom LoRA", visible=False)
         
         with gr.Column(scale=5):
             with gr.Row():
                 prompt = gr.Textbox(
-                    label="Editing Prompt",
+                    label="Additional Details (optional)",
                     show_label=False,
                     lines=1,
                     max_lines=1,
-                    placeholder="optional description, e.g. 'a man with glasses and a beard'",
+                    placeholder="Describe additional details, e.g., 'wearing a red hat' or 'smiling'",
                     elem_id="prompt"
                 )
                 run_button = gr.Button("Generate ✨", elem_id="run_button")
             
-            result = gr.Image(label="Generated Image", interactive=False)
+            result = gr.Image(label="Your Artistic Portrait", interactive=False)
             reuse_button = gr.Button("🔄 Reuse this image", visible=False)
             
-            with gr.Accordion("Advanced Settings", open=False):
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
                 lora_scale = gr.Slider(
-                    label="LoRA Scale",
+                    label="Style Strength",
                     minimum=0,
                     maximum=2,
                     step=0.1,
-                    value=1.5,
-                    info="Controls the strength of the LoRA effect"
+                    value=1.0,
+                    info="How strongly to apply the art style (1.0 = balanced)"
                 )
                 seed = gr.Slider(
-                    label="Seed",
+                    label="Random Seed",
                     minimum=0,
                     maximum=MAX_SEED,
                     step=1,
                     value=0,
+                    info="Set to 0 for random results"
                 )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                randomize_seed = gr.Checkbox(label="🎲 Randomize seed for each generation", value=True)
                 guidance_scale = gr.Slider(
-                    label="Guidance Scale",
+                    label="Image Guidance",
                     minimum=1,
                     maximum=10,
                     step=0.1,
                     value=2.5,
+                    info="How closely to follow the input image (lower = more creative)"
                 )
             
             prompt_title = gr.Markdown(
-                value="### Click on a LoRA in the gallery to select it",
+                value="### 🎨 Select an art style from the gallery",
                 visible=True,
                 elem_id="selected_lora",
             )