Update app.py

app.py

@@ -46,7 +46,7 @@ warnings.filterwarnings("ignore")
 # Set page config
 st.set_page_config(
     page_title="Style Transfer Studio",
-    # page_icon="🎨",
+    # page_icon="",
     layout="wide",
     initial_sidebar_state="expanded"
 )
@@ -1344,6 +1344,56 @@ class StyleTransferSystem:
                 return model
             except:
                 return None
+    # Inside the StyleTransferSystem class, add these methods:
+
+    def _create_linear_weight(self, width, height, overlap):
+        """Create linear blending weights for tile edges"""
+        weight = np.ones((height, width, 1), dtype=np.float32)
+        
+        if overlap > 0:
+            # Create gradients for each edge
+            for i in range(overlap):
+                alpha = i / overlap
+                # Top edge
+                weight[i, :] *= alpha
+                # Bottom edge
+                weight[-i-1, :] *= alpha
+                # Left edge
+                weight[:, i] *= alpha
+                # Right edge  
+                weight[:, -i-1] *= alpha
+        
+        return weight
+    
+    def _create_gaussian_weight(self, width, height, overlap):
+        """Create Gaussian blending weights for smoother transitions"""
+        weight = np.ones((height, width), dtype=np.float32)
+        
+        # Create 2D Gaussian centered in the tile
+        y, x = np.ogrid[:height, :width]
+        center_y, center_x = height / 2, width / 2
+        
+        # Distance from center
+        dist_from_center = np.sqrt((x - center_x)**2 + (y - center_y)**2)
+        
+        # Gaussian falloff starting from the edges
+        max_dist = min(height, width) / 2
+        sigma = max_dist / 2  # Adjust for smoother/sharper transitions
+        
+        # Apply Gaussian only near edges
+        edge_dist = np.minimum(
+            np.minimum(y, height - 1 - y),
+            np.minimum(x, width - 1 - x)
+        )
+        
+        # Weight is 1 in center, Gaussian falloff near edges
+        weight = np.where(
+            edge_dist < overlap,
+            np.exp(-0.5 * ((overlap - edge_dist) / (overlap/3))**2),
+            1.0
+        )
+        
+        return weight.reshape(height, width, 1)
 
     def apply_lightweight_style(self, image, model, intensity=1.0):
         """Apply style using a lightweight model"""
@@ -1654,8 +1704,18 @@ class StyleTransferSystem:
         
         return model
 
-    def apply_adain_style(self, content_image, style_image, model, alpha=1.0):
-        """Apply AdaIN-based style transfer with better quality preservation"""
+    def apply_adain_style(self, content_image, style_image, model, alpha=1.0, use_tiling=False):
+        """Apply AdaIN-based style transfer with optional tiling"""
+        if use_tiling and (content_image.width > 512 or content_image.height > 512):
+            # Use tiling for large images
+            return self.apply_adain_style_tiled(
+                content_image, style_image, model, alpha,
+                tile_size=256,  # Match training size
+                overlap=32,
+                blend_mode='gaussian'
+            )
+        
+        # Original implementation for small images
         if content_image is None or style_image is None or model is None:
             return None
         
@@ -1665,32 +1725,9 @@ class StyleTransferSystem:
             
             original_size = content_image.size
             
-            # Determine processing size based on GPU memory
-            if self.device.type == 'cuda':
-                # Try to use higher resolution on GPU
-                max_dimension = 768  # Increased from 256
-                
-                # Calculate memory-efficient size
-                aspect_ratio = content_image.width / content_image.height
-                if content_image.width > content_image.height:
-                    process_width = min(max_dimension, content_image.width)
-                    process_height = int(process_width / aspect_ratio)
-                else:
-                    process_height = min(max_dimension, content_image.height)
-                    process_width = int(process_height * aspect_ratio)
-                
-                # Round to multiples of 32 for better GPU efficiency
-                process_width = ((process_width + 31) // 32) * 32
-                process_height = ((process_height + 31) // 32) * 32
-            else:
-                # Lower resolution for CPU
-                process_width, process_height = 512, 512
-            
-            print(f"Processing at {process_width}x{process_height} (original: {original_size})")
-            
-            # Transform without aggressive cropping
+            # Transform for AdaIN (VGG normalization)
             transform = transforms.Compose([
-                transforms.Resize((process_height, process_width)),
+                transforms.Resize((256, 256)),  # Direct resize, no cropping
                 transforms.ToTensor(),
                 transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                      std=[0.229, 0.224, 0.225])
@@ -1700,71 +1737,117 @@ class StyleTransferSystem:
             style_tensor = transform(style_image).unsqueeze(0).to(self.device)
             
             with torch.no_grad():
-                # Process in chunks if image is very large
-                if process_width * process_height > 512 * 512 and self.device.type == 'cuda':
-                    # Clear cache before processing large image
-                    torch.cuda.empty_cache()
-                
                 output = model(content_tensor, style_tensor, alpha=alpha)
                 
-                # Improved denormalization
+                # Denormalize
                 output = output.squeeze(0).cpu()
-                
-                # More precise denormalization
-                denorm_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-                denorm_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
-                output = output * denorm_std + denorm_mean
-                
-                # Ensure values are in valid range
+                output = output * torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+                output = output + torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
                 output = torch.clamp(output, 0, 1)
                 
-                # Convert to PIL with better quality
+                # Convert to PIL
                 output_img = transforms.ToPILImage()(output)
-                
-                # Use LANCZOS for high-quality resizing
-                if output_img.size != original_size:
-                    output_img = output_img.resize(original_size, Image.LANCZOS)
-            
-            # Optional: Sharpen slightly to compensate for softness
-            if hasattr(Image, 'UnsharpMask'):
-                from PIL import ImageFilter
-                output_img = output_img.filter(ImageFilter.UnsharpMask(radius=1, percent=50, threshold=0))
+                output_img = output_img.resize(original_size, Image.LANCZOS)
             
             return output_img
         
-        except RuntimeError as e:
-            if "out of memory" in str(e):
-                print("GPU out of memory, falling back to lower resolution...")
-                torch.cuda.empty_cache()
-                
-                # Fallback to lower resolution
-                transform = transforms.Compose([
-                    transforms.Resize((384, 384)),  # Still better than 256
-                    transforms.ToTensor(),
-                    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                                         std=[0.229, 0.224, 0.225])
-                ])
-                
-                content_tensor = transform(content_image).unsqueeze(0).to(self.device)
-                style_tensor = transform(style_image).unsqueeze(0).to(self.device)
-                
-                with torch.no_grad():
-                    output = model(content_tensor, style_tensor, alpha=alpha)
-                    output = output.squeeze(0).cpu()
-                    
-                    denorm_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-                    denorm_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
-                    output = output * denorm_std + denorm_mean
-                    output = torch.clamp(output, 0, 1)
-                    
-                    output_img = transforms.ToPILImage()(output)
-                    output_img = output_img.resize(original_size, Image.LANCZOS)
-                
-                return output_img
-            else:
-                print(f"Error applying AdaIN style: {e}")
-                traceback.print_exc()
-                return None
+        except Exception as e:
+            print(f"Error applying AdaIN style: {e}")
+            traceback.print_exc()
+            return None
+
+    def apply_adain_style_tiled(self, content_image, style_image, model, alpha=1.0, 
+                            tile_size=256, overlap=32, blend_mode='linear'):
+        """
+        Apply AdaIN style transfer using tiling for high-quality results.
+        Processes image in overlapping tiles to maintain quality.
+        """
+        if content_image is None or style_image is None or model is None:
+            return None
+        
+        try:
+            model = model.to(self.device)
+            model.eval()
+            
+            # Prepare transforms
+            transform = transforms.Compose([
+                transforms.Resize((tile_size, tile_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+            
+            # Process style image once (at tile size)
+            style_tensor = transform(style_image).unsqueeze(0).to(self.device)
+            
+            # Get dimensions
+            w, h = content_image.size
+            
+            # Calculate tile positions with overlap
+            stride = tile_size - overlap
+            tiles_x = list(range(0, w - tile_size + 1, stride))
+            tiles_y = list(range(0, h - tile_size + 1, stride))
+            
+            # Ensure we cover the entire image
+            if tiles_x[-1] + tile_size < w:
+                tiles_x.append(w - tile_size)
+            if tiles_y[-1] + tile_size < h:
+                tiles_y.append(h - tile_size)
+            
+            # If image is smaller than tile size, just process normally
+            if w <= tile_size and h <= tile_size:
+                return self.apply_adain_style(content_image, style_image, model, alpha, use_tiling=False)
+            
+            print(f"Processing {len(tiles_x) * len(tiles_y)} tiles of size {tile_size}x{tile_size}")
+            
+            # Initialize output and weight arrays
+            output_array = np.zeros((h, w, 3), dtype=np.float32)
+            weight_array = np.zeros((h, w, 1), dtype=np.float32)
+            
+            # Process each tile
+            with torch.no_grad():
+                for y_idx, y in enumerate(tiles_y):
+                    for x_idx, x in enumerate(tiles_x):
+                        # Extract tile
+                        tile = content_image.crop((x, y, x + tile_size, y + tile_size))
+                        
+                        # Transform tile
+                        tile_tensor = transform(tile).unsqueeze(0).to(self.device)
+                        
+                        # Apply AdaIN to tile
+                        styled_tensor = model(tile_tensor, style_tensor, alpha=alpha)
+                        
+                        # Denormalize
+                        styled_tensor = styled_tensor.squeeze(0).cpu()
+                        denorm_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+                        denorm_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+                        styled_tensor = styled_tensor * denorm_std + denorm_mean
+                        styled_tensor = torch.clamp(styled_tensor, 0, 1)
+                        
+                        # Convert to numpy
+                        styled_tile = styled_tensor.permute(1, 2, 0).numpy() * 255
+                        
+                        # Create weight mask for blending
+                        if blend_mode == 'gaussian':
+                            weight = self._create_gaussian_weight(tile_size, tile_size, overlap)
+                        else:
+                            weight = self._create_linear_weight(tile_size, tile_size, overlap)
+                        
+                        # Add to output with weights
+                        output_array[y:y+tile_size, x:x+tile_size] += styled_tile * weight
+                        weight_array[y:y+tile_size, x:x+tile_size] += weight
+            
+            # Normalize by weights
+            output_array = output_array / (weight_array + 1e-8)
+            output_array = np.clip(output_array, 0, 255).astype(np.uint8)
+            
+            return Image.fromarray(output_array)
+            
+        except Exception as e:
+            print(f"Error in tiled AdaIN processing: {e}")
+            traceback.print_exc()
+            # Fallback to standard processing
+            return self.apply_adain_style(content_image, style_image, model, alpha, use_tiling=False)
 
             
 # ===========================
@@ -1801,6 +1884,64 @@ def combine_region_masks(canvas_results, canvas_size):
     
     return combined_mask
 
+def apply_adain_regional(content_image, style_image, model, canvas_result, alpha=1.0, feather_radius=10, use_tiling=False):
+    """Apply AdaIN style transfer to a painted region only"""
+    if content_image is None or style_image is None or model is None:
+        return None
+    
+    try:
+        # Get the mask from canvas
+        if canvas_result is None or canvas_result.image_data is None:
+            # No mask painted, apply to whole image
+            return system.apply_adain_style(content_image, style_image, model, alpha, use_tiling=use_tiling)
+        
+        # Extract mask from canvas
+        mask_data = canvas_result.image_data[:, :, 3]  # Alpha channel
+        mask = mask_data > 0
+        
+        # Resize mask to match original image size
+        original_size = content_image.size
+        display_size = (canvas_result.image_data.shape[1], canvas_result.image_data.shape[0])
+        
+        if original_size != display_size:
+            # Convert mask to PIL image for resizing
+            mask_pil = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
+            mask_pil = mask_pil.resize(original_size, Image.NEAREST)
+            mask = np.array(mask_pil) > 128
+        
+        # Apply feathering to mask edges if requested
+        if feather_radius > 0:
+            from scipy.ndimage import gaussian_filter
+            mask_float = mask.astype(np.float32)
+            mask_float = gaussian_filter(mask_float, sigma=feather_radius)
+            mask_float = np.clip(mask_float, 0, 1)
+        else:
+            mask_float = mask.astype(np.float32)
+        
+        # Apply style to entire image with tiling option
+        styled_full = system.apply_adain_style(content_image, style_image, model, alpha, use_tiling=use_tiling)
+        
+        if styled_full is None:
+            return None
+        
+        # Blend original and styled based on mask
+        original_array = np.array(content_image, dtype=np.float32)
+        styled_array = np.array(styled_full, dtype=np.float32)
+        
+        # Expand mask to 3 channels
+        mask_3ch = np.stack([mask_float] * 3, axis=2)
+        
+        # Blend
+        result_array = original_array * (1 - mask_3ch) + styled_array * mask_3ch
+        result_array = np.clip(result_array, 0, 255).astype(np.uint8)
+        
+        return Image.fromarray(result_array)
+        
+    except Exception as e:
+        print(f"Error applying regional AdaIN style: {e}")
+        traceback.print_exc()
+        return None
+
 # ===========================
 # INITIALIZE SYSTEM AND API
 # ===========================
@@ -1866,7 +2007,7 @@ with st.sidebar:
         st.caption("For faster processing, use a GPU-enabled environment")
     
     st.markdown("---")
-    st.markdown("### 📚 Quick Guide")
+    st.markdown("### Quick Guide")
     st.markdown("""
     - **Style Transfer**: Apply artistic styles to images
     - **Regional Transform**: Paint areas for local effects
@@ -2136,7 +2277,7 @@ with tab2:
             # Region management
             col_btn1, col_btn2, col_btn3 = st.columns(3)
             with col_btn1:
-                if st.button("➕ Add Region", use_container_width=True):
+                if st.button("Add Region", use_container_width=True):
                     new_region = {
                         'id': len(st.session_state.regions),
                         'style': style_choices[0] if style_choices else None,
@@ -2184,7 +2325,7 @@ with tab2:
                             key=f"region_intensity_{i}"
                         )
                     
-                    if st.button(f"🗑️ Remove Region {i+1}", key=f"remove_region_{i}"):
+                    if st.button(f"Remove Region {i+1}", key=f"remove_region_{i}"):
                         # Remove the region
                         st.session_state.regions.pop(i)
                         
@@ -2216,9 +2357,9 @@ with tab2:
             # Show workflow status
             if 'regional_result' in st.session_state:
                 if st.session_state.canvas_ready:
-                    st.success("**Edit Mode** - Paint your regions and click 'Apply Regional Styles' when ready")
+                    st.success("Edit Mode - Paint your regions and click 'Apply Regional Styles' when ready")
                 else:
-                    st.info("**Preview Mode** - Click 'Continue Editing' to modify regions")
+                    st.info("Preview Mode - Click 'Continue Editing' to modify regions")
             else:
                 st.info("Paint on the canvas below to define regions for each style")
             
@@ -2246,7 +2387,6 @@ with tab2:
             
             current_region = st.session_state.regions[current_region_idx]
             
-            # THIS IS THE FIX: The following line was added.
             col_draw1, col_draw2, col_draw3 = st.columns(3)
             
             with col_draw1:
@@ -2347,10 +2487,23 @@ with tab2:
                     
                     st.session_state['regional_result'] = result
         
-        # Show result
+        # Show result with fixed size
         if 'regional_result' in st.session_state:
             st.subheader("Result")
-            st.image(st.session_state['regional_result'], caption="Regional Styled Image", use_column_width=True)
+            
+            # Add display size control
+            display_size = st.slider("Display Size", 300, 800, 600, 50, key="regional_display_size")
+            
+            # Fixed size display
+            result_display = resize_image_for_display(
+                st.session_state['regional_result'], 
+                max_width=display_size, 
+                max_height=display_size
+            )
+            st.image(result_display, caption="Regional Styled Image")
+            
+            # Show actual dimensions
+            st.caption(f"Original size: {st.session_state['regional_result'].size[0]}x{st.session_state['regional_result'].size[1]} pixels")
             
             # Download button
             buf = io.BytesIO()
@@ -2361,10 +2514,7 @@ with tab2:
                 file_name=f"regional_styled_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
                 mime="image/png"
             )
-
-# TAB 3: Video Processing
-# TAB 3: Video Processing
-# TAB 3: Video Processing
+            
 # TAB 3: Video Processing
 with tab3:
     st.header("Video Processing")
@@ -2555,7 +2705,7 @@ with tab3:
 
 
 
-# TAB 4: Training with AdaIN
+# TAB 4: Training with AdaIN and Regional Application
 with tab4:
     st.header("Train Custom Style with AdaIN")
     st.markdown("Train your own style transfer model using Adaptive Instance Normalization")
@@ -2563,6 +2713,10 @@ with tab4:
     # Initialize session state for content images
     if 'content_images_list' not in st.session_state:
         st.session_state.content_images_list = []
+    if 'adain_canvas_result' not in st.session_state:
+        st.session_state.adain_canvas_result = None
+    if 'adain_test_image' not in st.session_state:
+        st.session_state.adain_test_image = None
     
     col1, col2, col3 = st.columns([1, 1, 1])
     
@@ -2603,7 +2757,7 @@ with tab4:
                 st.caption(f"... and {len(content_imgs) - 3} more")
     
     with col3:
-        st.subheader("⚙️ Training Settings")
+        st.subheader("Training Settings")
         
         model_name = st.text_input("Model Name", 
                                   value=f"adain_style_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}")
@@ -2664,18 +2818,24 @@ with tab4:
                         st.session_state['trained_adain_model'] = model
                         st.session_state['trained_style_images'] = style_images
                         st.session_state['model_path'] = f'/tmp/trained_models/{model_name}_final.pth'
-                        st.success("AdaIN training complete!")
+                        st.success("AdaIN training complete")
                     
                     progress_bar.empty()
                     status_text.empty()
             else:
                 st.error("Please upload both style and content images")
     
-    # Testing section
+    # Testing section with regional application
     if 'trained_adain_model' in st.session_state:
         st.markdown("---")
         st.header("Test Your AdaIN Model")
         
+        # Application mode selection
+        application_mode = st.radio("Application Mode", 
+                                   ["Whole Image", "Paint Region"], 
+                                   horizontal=True,
+                                   help="Choose whether to apply style to entire image or paint specific regions")
+        
         test_col1, test_col2, test_col3 = st.columns([1, 1, 1])
         
         with test_col1:
@@ -2683,7 +2843,7 @@ with tab4:
             
             # Test image selection
             test_source = st.radio("Test Image Source", 
-                                  ["Use Content Image", "Upload New"], 
+                                  ["Use Content Image", "Upload New", "Use Unsplash Image"], 
                                   horizontal=True)
             
             test_image = None
@@ -2693,6 +2853,13 @@ with tab4:
                                          range(len(st.session_state.content_images_list)),
                                          format_func=lambda x: f"Content Image {x+1}")
                 test_image = Image.open(st.session_state.content_images_list[content_idx]).convert('RGB')
+            elif test_source == "Use Unsplash Image":
+                # Use current Unsplash image if available
+                if 'current_image' in st.session_state and st.session_state.get('image_source', '').startswith('Unsplash'):
+                    test_image = st.session_state['current_image']
+                    st.success("Using Unsplash image")
+                else:
+                    st.info("Please search and select an image from the Style Transfer tab first")
             else:
                 # Upload new image
                 test_upload = st.file_uploader("Upload test image", 
@@ -2701,6 +2868,10 @@ with tab4:
                 if test_upload:
                     test_image = Image.open(test_upload).convert('RGB')
             
+            # Store test image in session state
+            if test_image:
+                st.session_state['adain_test_image'] = test_image
+            
             # Style selection for testing
             if 'trained_style_images' in st.session_state and len(st.session_state['trained_style_images']) > 1:
                 style_idx = st.selectbox("Select style", 
@@ -2716,36 +2887,125 @@ with tab4:
             # Alpha blending control
             alpha = st.slider("Style Strength (Alpha)", 0.0, 1.0, 1.0, 0.1,
                             help="0 = original content, 1 = full style transfer")
+            
+            # Add tiling option
+            use_tiling = st.checkbox("🔲 Use Tiled Processing", 
+                                    value=True,
+                                    help="Process large images in tiles for better quality. Recommended for images larger than 512x512.")
+            
+            # Regional painting options (only show if in paint mode)
+            if application_mode == "🖌️ Paint Region":
+                st.markdown("---")
+                st.subheader("🖌️ Painting Options")
+                
+                brush_size = st.slider("Brush Size", 5, 100, 30)
+                drawing_mode = st.selectbox("Drawing Tool", 
+                                          ["freedraw", "line", "rect", "circle", "polygon"],
+                                          index=0)
+                
+                # Feather/blur the mask edges
+                feather_radius = st.slider("Edge Softness", 0, 50, 10, 
+                                         help="Blur mask edges for smoother transitions")
+                
+                col_btn1, col_btn2 = st.columns(2)
+                with col_btn1:
+                    if st.button("Clear Canvas", use_container_width=True):
+                        st.session_state['adain_canvas_result'] = None
+                        st.rerun()
+                
+                with col_btn2:
+                    if st.button("Reset Result", use_container_width=True):
+                        if 'adain_styled_result' in st.session_state:
+                            del st.session_state['adain_styled_result']
+                        st.rerun()
         
         with test_col2:
-            st.subheader("Original")
-            if test_image:
-                st.image(test_image, caption="Content Image", use_column_width=True)
+            st.subheader("Canvas / Original")
+            
+            if application_mode == "Paint Region" and test_image:
+                # Show canvas for painting
+                display_img = resize_image_for_display(test_image, max_width=400, max_height=400)
+                canvas_width, canvas_height = display_img.size
+                
+                st.info("🖌️ Paint the areas where you want to apply the style")
+                
+                # Canvas for painting mask
+                canvas_result = st_canvas(
+                    fill_color="rgba(255, 0, 0, 0.3)",  # Red with transparency
+                    stroke_width=brush_size,
+                    stroke_color="rgba(255, 0, 0, 0.5)",
+                    background_image=display_img,
+                    update_streamlit=True,
+                    height=canvas_height,
+                    width=canvas_width,
+                    drawing_mode=drawing_mode,
+                    display_toolbar=True,
+                    key=f"adain_canvas_{brush_size}_{drawing_mode}"
+                )
+                
+                # Save canvas result
+                if canvas_result:
+                    st.session_state['adain_canvas_result'] = canvas_result
+                
+                # Show style image below canvas
                 if test_style:
+                    st.markdown("---")
                     st.image(test_style, caption="Style Image", use_column_width=True)
+            
+            else:
+                # Show original images
+                if test_image:
+                    st.image(test_image, caption="Content Image", use_column_width=True)
+                    if test_style:
+                        st.image(test_style, caption="Style Image", use_column_width=True)
         
         with test_col3:
             st.subheader("Result")
-            if test_image and test_style:
+            
+            # Apply button
+            apply_button = st.button("Apply Style", type="primary", use_container_width=True)
+            
+            if apply_button and test_image and test_style:
                 with st.spinner("Applying style..."):
-                    result = system.apply_adain_style(
-                        test_image, 
-                        test_style,
-                        st.session_state['trained_adain_model'],
-                        alpha=alpha
-                    )
-                    if result:
-                        st.image(result, caption="Styled Result", use_column_width=True)
-                        
-                        # Download button
-                        buf = io.BytesIO()
-                        result.save(buf, format='PNG')
-                        st.download_button(
-                            label="📥 Download Result",
-                            data=buf.getvalue(),
-                            file_name=f"adain_styled_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
-                            mime="image/png"
+                    if application_mode == "🖼️ Whole Image":
+                        # Apply to whole image
+                        result = system.apply_adain_style(
+                            test_image, 
+                            test_style,
+                            st.session_state['trained_adain_model'],
+                            alpha=alpha,
+                            use_tiling=use_tiling
+                        )
+                    else:
+                        # Apply to painted region
+                        result = apply_adain_regional(
+                            test_image,
+                            test_style,
+                            st.session_state['trained_adain_model'],
+                            st.session_state.get('adain_canvas_result'),
+                            alpha=alpha,
+                            feather_radius=feather_radius,
+                            use_tiling=use_tiling
                         )
+                    
+                    if result:
+                        st.session_state['adain_styled_result'] = result
+            
+            # Show result if available
+            if 'adain_styled_result' in st.session_state:
+                st.image(st.session_state['adain_styled_result'], 
+                        caption="Styled Result", 
+                        use_column_width=True)
+                
+                # Download button
+                buf = io.BytesIO()
+                st.session_state['adain_styled_result'].save(buf, format='PNG')
+                st.download_button(
+                    label="Download Result",
+                    data=buf.getvalue(),
+                    file_name=f"adain_styled_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
+                    mime="image/png"
+                )
         
         # Model download section
         st.markdown("---")
@@ -2756,13 +3016,73 @@ with tab4:
                     st.download_button(
                         label="Download Trained AdaIN Model",
                         data=f.read(),
-                        file_name=f"{model_name}_adain_final.pth",
+                        file_name=f"{st.session_state.get('model_name', 'adain')}_final.pth",
                         mime="application/octet-stream",
                         use_container_width=True
                     )
             with col_dl2:
                 st.info("This model can be loaded and used for real-time style transfer")
 
+
+# Add this helper function (place it before the tab or with other helper functions)
+def apply_adain_regional(content_image, style_image, model, canvas_result, alpha=1.0, feather_radius=10, use_tiling=False):
+    """Apply AdaIN style transfer to a painted region only"""
+    if content_image is None or style_image is None or model is None:
+        return None
+    
+    try:
+        # Get the mask from canvas
+        if canvas_result is None or canvas_result.image_data is None:
+            # No mask painted, apply to whole image
+            return system.apply_adain_style(content_image, style_image, model, alpha, use_tiling=use_tiling)
+        
+        # Extract mask from canvas
+        mask_data = canvas_result.image_data[:, :, 3]  # Alpha channel
+        mask = mask_data > 0
+        
+        # Resize mask to match original image size
+        original_size = content_image.size
+        display_size = (canvas_result.image_data.shape[1], canvas_result.image_data.shape[0])
+        
+        if original_size != display_size:
+            # Convert mask to PIL image for resizing
+            mask_pil = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
+            mask_pil = mask_pil.resize(original_size, Image.NEAREST)
+            mask = np.array(mask_pil) > 128
+        
+        # Apply feathering to mask edges if requested
+        if feather_radius > 0:
+            from scipy.ndimage import gaussian_filter
+            mask_float = mask.astype(np.float32)
+            mask_float = gaussian_filter(mask_float, sigma=feather_radius)
+            mask_float = np.clip(mask_float, 0, 1)
+        else:
+            mask_float = mask.astype(np.float32)
+        
+        # Apply style to entire image with tiling option
+        styled_full = system.apply_adain_style(content_image, style_image, model, alpha, use_tiling=use_tiling)
+        
+        if styled_full is None:
+            return None
+        
+        # Blend original and styled based on mask
+        original_array = np.array(content_image, dtype=np.float32)
+        styled_array = np.array(styled_full, dtype=np.float32)
+        
+        # Expand mask to 3 channels
+        mask_3ch = np.stack([mask_float] * 3, axis=2)
+        
+        # Blend
+        result_array = original_array * (1 - mask_3ch) + styled_array * mask_3ch
+        result_array = np.clip(result_array, 0, 255).astype(np.uint8)
+        
+        return Image.fromarray(result_array)
+        
+    except Exception as e:
+        print(f"Error applying regional AdaIN style: {e}")
+        traceback.print_exc()
+        return None
+        
 # TAB 5: Batch Processing
 with tab5:
     st.header("Batch Processing")
@@ -2915,7 +3235,7 @@ with tab6:
     - Supports all style combinations and blend modes
     - Enhanced codec compatibility
     
-    #### 🔧 Custom Training
+    #### Custom Training
     - Train on any artistic style with minimal data (1-50 images)
     - Automatic data augmentation for small datasets
     - Adjustable model complexity (3-12 residual blocks)
@@ -2976,4 +3296,4 @@ with tab6:
 
 # Footer
 st.markdown("---")
-st.markdown("Style transfer system with state-of-the-art CycleGAN models and regional painting capabilities.")
+st.markdown("Style transfer system with CycleGAN models and regional painting capabilities.")