Update app.py

app.py

@@ -181,40 +181,32 @@ def style_transfer(input_image, style_name, prompt_suffix, num_inference_steps,
         torch.cuda.empty_cache()
         return None
 
-def create_thumbnail_html():
-    """Create HTML for thumbnail grid"""
-    html = '<div style="display: grid; grid-template-columns: repeat(6, 1fr); gap: 10px; max-width: 800px; margin: 0 auto;">'
-    
+def select_style(style_name):
+    """Handler for thumbnail clicks"""
+    return style_name, style_descriptions.get(style_name, "")
+
+def create_thumbnail_grid():
+    """Create a gallery of style thumbnails"""
+    thumbnails = []
     styles = list(style_type_lora_dict.keys())
     
-    for i, style in enumerate(styles):
-        if i >= 24:  # Limit to 24 thumbnails for 6x4 grid
-            break
-        
+    for style in styles:
         thumbnail_file = thumbnail_mapping.get(style, "")
-        style_readable = style.replace('_', ' ')
-        
-        html += f'''
-        <div style="text-align: center; cursor: pointer;" onclick="document.getElementById('style_dropdown').value='{style}'; 
-                     var event = new Event('change', {{bubbles: true}});
-                     document.getElementById('style_dropdown').dispatchEvent(event);">
-            <img src="file/{thumbnail_file}" alt="{style_readable}" 
-                 style="width: 100%; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); 
-                        transition: transform 0.2s, box-shadow 0.2s;"
-                 onmouseover="this.style.transform='scale(1.05)'; this.style.boxShadow='0 4px 8px rgba(0,0,0,0.2)';"
-                 onmouseout="this.style.transform='scale(1)'; this.style.boxShadow='0 2px 4px rgba(0,0,0,0.1)';">
-            <p style="margin: 5px 0; font-size: 12px; font-weight: 500;">{style_readable}</p>
-        </div>
-        '''
-    
-    # Fill empty slots if needed
-    remaining_slots = 24 - len(styles)
-    if remaining_slots > 0 and len(styles) < 24:
-        for _ in range(remaining_slots):
-            html += '<div></div>'
+        if thumbnail_file and os.path.exists(thumbnail_file):
+            try:
+                img = Image.open(thumbnail_file)
+                thumbnails.append((img, style.replace('_', ' ')))
+            except Exception as e:
+                print(f"Error loading thumbnail {thumbnail_file}: {e}")
+                # Create placeholder if thumbnail fails to load
+                placeholder = Image.new('RGB', (256, 256), color='lightgray')
+                thumbnails.append((placeholder, style.replace('_', ' ')))
+        else:
+            # Create placeholder for missing thumbnails
+            placeholder = Image.new('RGB', (256, 256), color='lightgray')
+            thumbnails.append((placeholder, style.replace('_', ' ')))
     
-    html += '</div>'
-    return html
+    return thumbnails
 
 # Create Gradio interface
 with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as demo:
@@ -228,8 +220,20 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
     
     # Thumbnail Grid Section
     gr.Markdown("### 🖼️ Click a style thumbnail to select it:")
+    
     with gr.Row():
-        gr.HTML(create_thumbnail_html())
+        style_gallery = gr.Gallery(
+            value=create_thumbnail_grid(),
+            label="Style Thumbnails",
+            show_label=False,
+            elem_id="style_gallery",
+            columns=6,
+            rows=4,
+            object_fit="cover",
+            height="auto",
+            interactive=True,
+            show_download_button=False
+        )
     
     gr.Markdown("---")
     
@@ -307,6 +311,22 @@ with gr.Blocks(title="FLUX.1 Kontext Style Transfer", theme=gr.themes.Soft()) as
             - Use additional instructions for fine control
             """)
     
+    # Handle gallery selection
+    def on_gallery_select(evt: gr.SelectData):
+        """Handle thumbnail selection from gallery"""
+        selected_index = evt.index
+        styles = list(style_type_lora_dict.keys())
+        if 0 <= selected_index < len(styles):
+            selected_style = styles[selected_index]
+            return selected_style, style_descriptions.get(selected_style, "")
+        return None, None
+    
+    style_gallery.select(
+        fn=on_gallery_select,
+        inputs=None,
+        outputs=[style_dropdown, style_info]
+    )
+    
     # Update style description when style changes
     def update_description(style):
         return style_descriptions.get(style, "")