Update app.py

app.py

@@ -1570,14 +1570,17 @@ class StyleTransferSystem:
         return Image.fromarray(np.clip(result, 0, 255).astype(np.uint8))
 
     def train_adain_model(self, style_images, content_dir, model_name,
-                          epochs=30, batch_size=4, lr=1e-4,
-                          save_interval=5, style_weight=10.0, content_weight=1.0,
-                          progress_callback=None):
+                      epochs=30, batch_size=4, lr=1e-4,
+                      save_interval=5, style_weight=10.0, content_weight=1.0,
+                      progress_callback=None):
         """Train an AdaIN-based style transfer model"""
         
         model = AdaINStyleTransfer().to(self.device)
         optimizer = torch.optim.Adam(model.decoder.parameters(), lr=lr)
         
+        # Add learning rate scheduler
+        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.8)
+        
         print(f"Training AdaIN model")
         print(f"Training device: {self.device}")
         
@@ -1586,24 +1589,32 @@ class StyleTransferSystem:
             print(f"Model on GPU: {next(model.decoder.parameters()).device}")
             print(f"GPU memory before training: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
         
-        # Prepare style images
+        # Prepare style images - INCREASED SIZE
         style_transform = transforms.Compose([
-            transforms.Resize(512),
-            transforms.RandomCrop(256),
+            transforms.Resize(600),  # Increased from 512
+            transforms.RandomCrop(512),  # Increased from 256
+            transforms.RandomHorizontalFlip(p=0.5),  # Add augmentation
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225])
         ])
         
         style_tensors = []
+        # Create multiple augmented versions of each style image
         for style_img in style_images:
-            style_tensor = style_transform(style_img).unsqueeze(0).to(self.device)
-            style_tensors.append(style_tensor)
+            # Generate 5 augmented versions per style image
+            for _ in range(5):
+                style_tensor = style_transform(style_img).unsqueeze(0).to(self.device)
+                style_tensors.append(style_tensor)
+        
+        print(f"Created {len(style_tensors)} augmented style samples from {len(style_images)} images")
         
-        # Prepare content dataset
+        # Prepare content dataset - INCREASED SIZE
         content_transform = transforms.Compose([
-            transforms.Resize(512),
-            transforms.RandomCrop(256),
+            transforms.Resize(600),  # Increased from 512
+            transforms.RandomCrop(512),  # Increased from 256
+            transforms.RandomHorizontalFlip(),
+            transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225])
@@ -1617,9 +1628,28 @@ class StyleTransferSystem:
         print(f"  - Content images: {len(dataset)}")
         print(f"  - Batch size: {batch_size}")
         print(f"  - Epochs: {epochs}")
-        
-        # Loss network (VGG for perceptual loss)
-        loss_network = VGGEncoder().to(self.device).eval()
+        print(f"  - Training resolution: 512x512")  # Updated
+        
+        # Loss network (VGG for perceptual loss) - USE MULTIPLE LAYERS
+        class MultiLayerVGG(nn.Module):
+            def __init__(self):
+                super().__init__()
+                vgg = models.vgg19(weights=models.VGG19_Weights.DEFAULT).features
+                self.slice1 = nn.Sequential(*list(vgg.children())[:2])    # relu1_1
+                self.slice2 = nn.Sequential(*list(vgg.children())[2:7])   # relu2_1
+                self.slice3 = nn.Sequential(*list(vgg.children())[7:12])  # relu3_1
+                self.slice4 = nn.Sequential(*list(vgg.children())[12:21]) # relu4_1
+                for param in self.parameters():
+                    param.requires_grad = False
+            
+            def forward(self, x):
+                h1 = self.slice1(x)
+                h2 = self.slice2(h1)
+                h3 = self.slice3(h2)
+                h4 = self.slice4(h3)
+                return [h1, h2, h3, h4]
+        
+        loss_network = MultiLayerVGG().to(self.device).eval()
         mse_loss = nn.MSELoss()
         
         # Training loop
@@ -1627,6 +1657,9 @@ class StyleTransferSystem:
         model.encoder.eval()  # Keep encoder frozen
         total_steps = 0
         
+        # Adjust style weight for better quality
+        actual_style_weight = style_weight * 10  # Multiply by 10 for better style transfer
+        
         for epoch in range(epochs):
             epoch_loss = 0
             
@@ -1643,33 +1676,42 @@ class StyleTransferSystem:
                 # Forward pass
                 output = model(content_batch, batch_style)
                 
-                # Content loss
+                # Multi-layer content and style loss
                 with torch.no_grad():
-                    content_feat = loss_network.encode(content_batch)
-                output_feat = loss_network.encode(output)
-                content_loss = mse_loss(output_feat, content_feat)
+                    content_feats = loss_network(content_batch)
+                    style_feats = loss_network(batch_style)
+                output_feats = loss_network(output)
                 
-                # Style loss
-                with torch.no_grad():
-                    style_feat = loss_network.encode(batch_style)
+                # Content loss - only from relu4_1
+                content_loss = mse_loss(output_feats[-1], content_feats[-1])
+                
+                # Style loss - from multiple layers
+                style_loss = 0
+                style_weights = [0.2, 0.3, 0.5, 1.0]  # Give more weight to higher layers
                 
-                # Compute style loss using Gram matrices
                 def gram_matrix(feat):
                     b, c, h, w = feat.size()
                     feat = feat.view(b, c, h * w)
                     gram = torch.bmm(feat, feat.transpose(1, 2))
                     return gram / (c * h * w)
                 
-                output_gram = gram_matrix(output_feat)
-                style_gram = gram_matrix(style_feat)
-                style_loss = mse_loss(output_gram, style_gram)
+                for i, (output_feat, style_feat, weight) in enumerate(zip(output_feats, style_feats, style_weights)):
+                    output_gram = gram_matrix(output_feat)
+                    style_gram = gram_matrix(style_feat)
+                    style_loss += weight * mse_loss(output_gram, style_gram)
+                
+                style_loss /= len(style_weights)
                 
                 # Total loss
-                loss = content_weight * content_loss + style_weight * style_loss
+                loss = content_weight * content_loss + actual_style_weight * style_loss
                 
                 # Backward pass
                 optimizer.zero_grad()
                 loss.backward()
+                
+                # Gradient clipping for stability
+                torch.nn.utils.clip_grad_norm_(model.decoder.parameters(), max_norm=5.0)
+                
                 optimizer.step()
                 
                 epoch_loss += loss.item()
@@ -1678,7 +1720,12 @@ class StyleTransferSystem:
                 # Progress callback
                 if progress_callback and total_steps % 10 == 0:
                     progress = (epoch + (batch_idx + 1) / len(dataloader)) / epochs
-                    progress_callback(progress, f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}")
+                    progress_callback(progress, 
+                        f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f} "
+                        f"(Content: {content_loss.item():.4f}, Style: {style_loss.item():.4f})")
+            
+            # Step scheduler
+            scheduler.step()
             
             # Save checkpoint
             if (epoch + 1) % save_interval == 0:
@@ -1687,6 +1734,7 @@ class StyleTransferSystem:
                     'epoch': epoch + 1,
                     'model_state_dict': model.state_dict(),
                     'optimizer_state_dict': optimizer.state_dict(),
+                    'scheduler_state_dict': scheduler.state_dict(),
                     'loss': epoch_loss / len(dataloader),
                     'model_type': 'adain'
                 }, checkpoint_path)
@@ -1704,19 +1752,21 @@ class StyleTransferSystem:
         self.lightweight_models[model_name] = model
         
         return model
+        
+
+    # Update these methods in your StyleTransferSystem class:
 
     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
+        # Use tiling for large images to maintain quality
+        if use_tiling and (content_image.width > 768 or content_image.height > 768):
             return self.apply_adain_style_tiled(
                 content_image, style_image, model, alpha,
-                tile_size=256,  # Match training size
-                overlap=32,
+                tile_size=512,  # Increased from 256
+                overlap=64,     # Increased overlap
                 blend_mode='gaussian'
             )
         
-        # Original implementation for small images
         if content_image is None or style_image is None or model is None:
             return None
         
@@ -1726,9 +1776,23 @@ class StyleTransferSystem:
             
             original_size = content_image.size
             
+            # Use higher resolution - find optimal size while maintaining aspect ratio
+            max_dim = 768  # Increased from 256
+            w, h = content_image.size
+            if w > h:
+                new_w = min(w, max_dim)
+                new_h = int(h * new_w / w)
+            else:
+                new_h = min(h, max_dim)
+                new_w = int(w * new_h / h)
+            
+            # Ensure dimensions are divisible by 8 for better compatibility
+            new_w = (new_w // 8) * 8
+            new_h = (new_h // 8) * 8
+            
             # Transform for AdaIN (VGG normalization)
             transform = transforms.Compose([
-                transforms.Resize((256, 256)),  # Direct resize, no cropping
+                transforms.Resize((new_h, new_w)),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                      std=[0.229, 0.224, 0.225])
@@ -1756,7 +1820,7 @@ class StyleTransferSystem:
             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'):
         """
@@ -1770,6 +1834,10 @@ class StyleTransferSystem:
             model = model.to(self.device)
             model.eval()
             
+            # INCREASED TILE SIZE FOR BETTER QUALITY
+            tile_size = 512  # Override input to use 512
+            overlap = 64    # Increase overlap proportionally
+            
             # Prepare transforms
             transform = transforms.Compose([
                 transforms.Resize((tile_size, tile_size)),
@@ -1790,10 +1858,10 @@ class StyleTransferSystem:
             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 not tiles_x or tiles_x[-1] + tile_size < w:
+                tiles_x.append(max(0, w - tile_size))
+            if not tiles_y or tiles_y[-1] + tile_size < h:
+                tiles_y.append(max(0, h - tile_size))
             
             # If image is smaller than tile size, just process normally
             if w <= tile_size and h <= tile_size:
@@ -1828,11 +1896,8 @@ class StyleTransferSystem:
                         # 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)
+                        # Create weight mask for blending - use gaussian by default for better quality
+                        weight = self._create_gaussian_weight(tile_size, tile_size, overlap)
                         
                         # Add to output with weights
                         output_array[y:y+tile_size, x:x+tile_size] += styled_tile * weight
@@ -2706,6 +2771,7 @@ with tab3:
 
 
 
+# TAB 4: Training with AdaIN and Regional Application
 # TAB 4: Training with AdaIN and Regional Application
 with tab4:
     st.header("Train Custom Style with AdaIN")
@@ -2739,7 +2805,7 @@ with tab4:
     
     with col2:
         st.subheader("Content Images")
-        content_imgs = st.file_uploader("Upload content images (5-50 recommended)", 
+        content_imgs = st.file_uploader("Upload content images (10-50 recommended)", 
                                        type=['png', 'jpg', 'jpeg'], 
                                        accept_multiple_files=True, 
                                        key="train_content_adain")
@@ -2763,22 +2829,26 @@ with tab4:
         model_name = st.text_input("Model Name", 
                                   value=f"adain_style_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}")
         
-        epochs = st.slider("Training Epochs", 5, 50, 15, 5)
+        # IMPROVED DEFAULT VALUES
+        epochs = st.slider("Training Epochs", 10, 100, 50, 5)  # Increased default
         batch_size = st.slider("Batch Size", 1, 8, 4)
         learning_rate = st.number_input("Learning Rate", 0.00001, 0.001, 0.0001, format="%.5f")
         
         with st.expander("Advanced Settings"):
-            style_weight = st.number_input("Style Weight", 1.0, 100.0, 10.0, 1.0)
+            # MUCH HIGHER STYLE WEIGHT BY DEFAULT
+            style_weight = st.number_input("Style Weight", 1.0, 1000.0, 100.0, 10.0)
             content_weight = st.number_input("Content Weight", 0.1, 10.0, 1.0, 0.1)
-            save_interval = st.slider("Save Checkpoint Every N Epochs", 5, 20, 5, 5)
+            save_interval = st.slider("Save Checkpoint Every N Epochs", 5, 20, 10, 5)
+            
+            st.info("💡 **Pro tip**: For better quality, use Style Weight 100-500x higher than Content Weight")
         
         st.markdown("---")
         
         # Training button
         if st.button("Start AdaIN Training", type="primary", use_container_width=True):
             if style_imgs and content_imgs:
-                if len(content_imgs) < 5:
-                    st.warning("For best results, use at least 5 content images")
+                if len(content_imgs) < 10:
+                    st.warning("For best results, use at least 10 content images")
                 
                 with st.spinner("Training AdaIN model..."):
                     progress_bar = st.progress(0)
@@ -2803,14 +2873,170 @@ with tab4:
                         style_img = Image.open(style_file).convert('RGB')
                         style_images.append(style_img)
                     
-                    # Train model
-                    model = system.train_adain_model(
-                        style_images, temp_content_dir, model_name,
-                        epochs=epochs, lr=learning_rate, batch_size=batch_size,
-                        save_interval=save_interval, style_weight=style_weight,
-                        content_weight=content_weight,
-                        progress_callback=progress_callback
-                    )
+                    # IMPROVED TRAINING FUNCTION
+                    # Multi-layer VGG loss for better quality
+                    class MultiLayerVGG(nn.Module):
+                        def __init__(self):
+                            super().__init__()
+                            vgg = models.vgg19(weights=models.VGG19_Weights.DEFAULT).features
+                            self.slice1 = nn.Sequential(*list(vgg.children())[:2])    # relu1_1
+                            self.slice2 = nn.Sequential(*list(vgg.children())[2:7])   # relu2_1
+                            self.slice3 = nn.Sequential(*list(vgg.children())[7:12])  # relu3_1
+                            self.slice4 = nn.Sequential(*list(vgg.children())[12:21]) # relu4_1
+                            for param in self.parameters():
+                                param.requires_grad = False
+                        
+                        def forward(self, x):
+                            h1 = self.slice1(x)
+                            h2 = self.slice2(h1)
+                            h3 = self.slice3(h2)
+                            h4 = self.slice4(h3)
+                            return [h1, h2, h3, h4]
+                    
+                    # Create model
+                    model = AdaINStyleTransfer().to(system.device)
+                    optimizer = torch.optim.Adam(model.decoder.parameters(), lr=learning_rate)
+                    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.8)
+                    
+                    print(f"Training AdaIN model at 512x512 resolution")
+                    print(f"Training device: {system.device}")
+                    
+                    # Prepare style images - LARGER SIZE
+                    style_transform = transforms.Compose([
+                        transforms.Resize(600),  # Increased size
+                        transforms.RandomCrop(512),  # Larger crops
+                        transforms.RandomHorizontalFlip(p=0.5),
+                        transforms.ToTensor(),
+                        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                             std=[0.229, 0.224, 0.225])
+                    ])
+                    
+                    style_tensors = []
+                    # Create multiple augmented versions
+                    for style_img in style_images:
+                        for _ in range(5):  # 5 augmented versions per style
+                            style_tensor = style_transform(style_img).unsqueeze(0).to(system.device)
+                            style_tensors.append(style_tensor)
+                    
+                    # Prepare content dataset - LARGER SIZE
+                    content_transform = transforms.Compose([
+                        transforms.Resize(600),
+                        transforms.RandomCrop(512),
+                        transforms.RandomHorizontalFlip(),
+                        transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05),
+                        transforms.ToTensor(),
+                        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                             std=[0.229, 0.224, 0.225])
+                    ])
+                    
+                    dataset = StyleTransferDataset(temp_content_dir, transform=content_transform)
+                    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)
+                    
+                    # Multi-layer loss network
+                    loss_network = MultiLayerVGG().to(system.device).eval()
+                    mse_loss = nn.MSELoss()
+                    
+                    # Training loop
+                    model.train()
+                    model.encoder.eval()
+                    total_steps = 0
+                    
+                    # Multiply style weight for better results
+                    actual_style_weight = style_weight * 10
+                    
+                    for epoch in range(epochs):
+                        epoch_loss = 0
+                        epoch_content_loss = 0
+                        epoch_style_loss = 0
+                        
+                        for batch_idx, content_batch in enumerate(dataloader):
+                            content_batch = content_batch.to(system.device)
+                            
+                            # Randomly select style images
+                            batch_style = []
+                            for _ in range(content_batch.size(0)):
+                                style_idx = np.random.randint(0, len(style_tensors))
+                                batch_style.append(style_tensors[style_idx])
+                            batch_style = torch.cat(batch_style, dim=0)
+                            
+                            # Forward pass
+                            output = model(content_batch, batch_style)
+                            
+                            # Multi-layer loss
+                            with torch.no_grad():
+                                content_feats = loss_network(content_batch)
+                                style_feats = loss_network(batch_style)
+                            output_feats = loss_network(output)
+                            
+                            # Content loss from relu4_1
+                            content_loss = mse_loss(output_feats[-1], content_feats[-1])
+                            
+                            # Style loss from multiple layers
+                            style_loss = 0
+                            style_weights = [0.2, 0.3, 0.5, 1.0]
+                            
+                            def gram_matrix(feat):
+                                b, c, h, w = feat.size()
+                                feat = feat.view(b, c, h * w)
+                                gram = torch.bmm(feat, feat.transpose(1, 2))
+                                return gram / (c * h * w)
+                            
+                            for i, (output_feat, style_feat, weight) in enumerate(zip(output_feats, style_feats, style_weights)):
+                                output_gram = gram_matrix(output_feat)
+                                style_gram = gram_matrix(style_feat)
+                                style_loss += weight * mse_loss(output_gram, style_gram)
+                            
+                            style_loss /= len(style_weights)
+                            
+                            # Total loss
+                            loss = content_weight * content_loss + actual_style_weight * style_loss
+                            
+                            # Backward pass
+                            optimizer.zero_grad()
+                            loss.backward()
+                            torch.nn.utils.clip_grad_norm_(model.decoder.parameters(), max_norm=5.0)
+                            optimizer.step()
+                            
+                            epoch_loss += loss.item()
+                            epoch_content_loss += content_loss.item()
+                            epoch_style_loss += style_loss.item()
+                            total_steps += 1
+                            
+                            # Progress callback
+                            if progress_callback and total_steps % 10 == 0:
+                                progress = (epoch + (batch_idx + 1) / len(dataloader)) / epochs
+                                progress_callback(progress, 
+                                    f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f} "
+                                    f"(C: {content_loss.item():.4f}, S: {style_loss.item():.4f})")
+                        
+                        # Step scheduler
+                        scheduler.step()
+                        
+                        # Print epoch stats
+                        avg_loss = epoch_loss / len(dataloader)
+                        print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, "
+                              f"Content={epoch_content_loss/len(dataloader):.4f}, "
+                              f"Style={epoch_style_loss/len(dataloader):.4f}")
+                        
+                        # Save checkpoint
+                        if (epoch + 1) % save_interval == 0:
+                            checkpoint_path = f'{system.models_dir}/{model_name}_epoch_{epoch+1}.pth'
+                            torch.save({
+                                'epoch': epoch + 1,
+                                'model_state_dict': model.state_dict(),
+                                'optimizer_state_dict': optimizer.state_dict(),
+                                'scheduler_state_dict': scheduler.state_dict(),
+                                'loss': avg_loss,
+                                'model_type': 'adain'
+                            }, checkpoint_path)
+                            print(f"Saved checkpoint: {checkpoint_path}")
+                    
+                    # Save final model
+                    final_path = f'{system.models_dir}/{model_name}_final.pth'
+                    torch.save({
+                        'model_state_dict': model.state_dict(),
+                        'model_type': 'adain'
+                    }, final_path)
                     
                     # Cleanup
                     shutil.rmtree(temp_content_dir)
@@ -2818,8 +3044,11 @@ with tab4:
                     if model:
                         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.session_state['model_path'] = final_path
+                        st.success("AdaIN training complete! 🎉")
+                        
+                        # Add to system's models
+                        system.lightweight_models[model_name] = model
                     
                     progress_bar.empty()
                     status_text.empty()
@@ -2885,14 +3114,15 @@ with tab4:
             else:
                 test_style = None
             
+            # IMPROVED DEFAULTS
             # 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")
+            alpha = st.slider("Style Strength (Alpha)", 0.0, 2.0, 1.2, 0.1,
+                            help="0 = original content, 1 = full style transfer, >1 = stronger style")
             
-            # Add tiling option
+            # Add tiling option - DEFAULT TO TRUE
             use_tiling = st.checkbox("Use Tiled Processing", 
-                                    value=True,
-                                    help="Process large images in tiles for better quality. Recommended for images larger than 512x512.")
+                                    value=True,  # Default to True
+                                    help="Process images in tiles for better quality. Recommended for ALL images.")
             
             # Initialize variables with default values
             brush_size = 30
@@ -3003,6 +3233,16 @@ with tab4:
                         caption="Styled Result", 
                         use_column_width=True)
                 
+                # Quality tips
+                with st.expander("💡 Tips for Better Quality"):
+                    st.markdown("""
+                    - **Always use tiling** for best quality
+                    - Try **alpha > 1.0** (1.2-1.5) for stronger style
+                    - Use **multiple style images** when training
+                    - Train for **50+ epochs** for best results
+                    - If quality is still poor, retrain with **style weight = 200-500**
+                    """)
+                
                 # Download button
                 buf = io.BytesIO()
                 st.session_state['adain_styled_result'].save(buf, format='PNG')
@@ -3022,7 +3262,7 @@ with tab4:
                     st.download_button(
                         label="Download Trained AdaIN Model",
                         data=f.read(),
-                        file_name=f"{st.session_state.get('model_name', 'adain')}_final.pth",
+                        file_name=f"{model_name}_final.pth",
                         mime="application/octet-stream",
                         use_container_width=True
                     )