cleaning

app.py

@@ -4,34 +4,27 @@ import numpy as np
 import gradio as gr
 import os
 
-from model import ChaoticCoherentGenerator
-# from model_256 import EfficientChaoticGenerator
+from model import Generator
 from feature_extractor import CodeFeatureExtractor
 
-# ------------------- Device -------------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ------------------- Load Model -------------------
-model = ChaoticCoherentGenerator().to(device)
-checkpoint = torch.load("models/chotic.pth", map_location=device)
+model = Generator().to(device)
+checkpoint = torch.load("models/mode_32.pth", map_location=device)
 model.load_state_dict(checkpoint["state_dict"])
 model.eval()
 
 extractor = CodeFeatureExtractor()
 
-# ------------------- Image Enhancement -------------------
 def enhance_image(image: Image.Image, upscale_size: int):
-    # Upscale smoothly
     image = image.resize((upscale_size, upscale_size), Image.Resampling.BICUBIC)
     
-    # Optional post-processing
     image = image.filter(ImageFilter.GaussianBlur(radius=0.8))
     image = ImageEnhance.Color(image).enhance(1.2)
     image = ImageEnhance.Sharpness(image).enhance(1.1)
     
     return image
 
-# ------------------- Generation Function -------------------
 def generate_from_code(code_text, upscale_size):
     temp_file = "temp.py"
     with open(temp_file, "w", encoding="utf-8") as f:
@@ -51,7 +44,6 @@ def generate_from_code(code_text, upscale_size):
     enhanced = enhance_image(img, upscale_size)
     return enhanced
 
-# ------------------- Gradio UI -------------------
 demo = gr.Interface(
     fn=generate_from_code,
     inputs=[

feature_extractor.py

@@ -38,15 +38,6 @@ class CodeFeatureExtractor:
         ]
     
     def extract_from_file(self, filepath):
-        """
-        Extract features from file
-        
-        Args:
-            filepath (str)
-            
-        Returns:
-            list: list of dfeatures
-        """
         try:
             with open(filepath, 'r', encoding='utf-8') as f:
                 code = f.read()
@@ -229,13 +220,6 @@ class CodeFeatureExtractor:
         return volume
     
     def extract_from_directory(self, directory, output_file='data/features.npy'):
-        """
-        Extract from a dir
-        
-        Args:
-            directory (str): dir name
-            output_file (str): file to save
-        """
         print(f"Extracting from: {directory}")
         
         features_list = []
@@ -258,7 +242,7 @@ class CodeFeatureExtractor:
                 failed += 1
         
         print(f"\n Features extracted from {len(features_list)} files")
-        print(f"⚠️Failed: {failed} files")
+        print(f"Failed: {failed} files")
         
         features_array = np.array(features_list)
         np.save(output_file, features_array)
@@ -290,10 +274,6 @@ class CodeFeatureExtractor:
 def main():
     extractor = CodeFeatureExtractor()
     
-    print("="*60)
-    print("Feature Extractor Code")
-    print("="*60)
-    
     #features = extractor.extract_from_directory('data/raw_code')
     
     features = extractor.extract_from_file('src/model.py')

model.py

@@ -7,7 +7,6 @@ from torch.utils.data import Dataset, DataLoader
 import matplotlib.pyplot as plt
 
 
-# ===================== Feature Names =====================
 
 FEATURE_NAMES = [
     "lines_of_code", "num_functions", "num_classes", "num_loops",
@@ -20,7 +19,6 @@ FEATURE_NAMES = [
 ]
 
 
-# ===================== Dataset =====================
 
 class CodeToImageDataset(Dataset):
     def __init__(self, features):
@@ -33,28 +31,27 @@ class CodeToImageDataset(Dataset):
         return self.features[idx]
 
 
-# ===================== Feature-Driven Generator with CNN =====================
 
-class ChaoticCoherentGenerator(nn.Module):
+class Generator(nn.Module):
     """
-    Generates images with controlled chaos through:
-    1. Frequency-domain decomposition (low/mid/high)
-    2. Spatial composition layers
-    3. Feature-dependent color palettes
-    4. Texture injection at multiple scales
+    generates images with controlled chaos through:
+     Frequency-domain decomposition (low/mid/high)
+     Spatial composition layers
+     Feature-dependent color palettes
+     Texture injection at multiple scales
     """
     def __init__(self, input_dim=25, image_size=32):
         super().__init__()
         self.image_size = image_size
 
-        # Feature grouping (keep your existing groups)
+        # Feature groupin
         self.structure_idx = [0, 1, 2, 10, 19, 20]
         self.control_idx = [3, 4, 9]
         self.operations_idx = [5, 7, 8, 6]
         self.style_idx = [14, 15, 24, 23]
         self.advanced_idx = [12, 13, 16, 17, 18, 11, 21, 22]
 
-        # Low-freq: overall composition, large shapes (structure features)
+        #overall composition, large shapes (structure features)
         self.low_freq_encoder = nn.Sequential(
             nn.Linear(len(self.structure_idx), 256),
             nn.LayerNorm(256),
@@ -63,7 +60,7 @@ class ChaoticCoherentGenerator(nn.Module):
             nn.Linear(256, 512)
         )
 
-        # Mid-freq: patterns, textures (control + operations)
+        # patterns, textures (control + operations)
         self.mid_freq_encoder = nn.Sequential(
             nn.Linear(len(self.control_idx) + len(self.operations_idx), 256),
             nn.LayerNorm(256),
@@ -72,7 +69,7 @@ class ChaoticCoherentGenerator(nn.Module):
             nn.Linear(256, 512)
         )
 
-        # High-freq: details, noise (style + advanced)
+        # details, noise (style + advanced)
         self.high_freq_encoder = nn.Sequential(
             nn.Linear(len(self.style_idx) + len(self.advanced_idx), 256),
             nn.LayerNorm(256),
@@ -88,7 +85,7 @@ class ChaoticCoherentGenerator(nn.Module):
             nn.Linear(128, 12)  # 4 colors × 3 channels (HSV or RGB)
         )
 
-        # Controls WHERE patterns appear (position, scale, rotation)
+        # WHERE patterns appear (position, scale, rotation)
         self.composition_net = nn.Sequential(
             nn.Linear(512, 256),
             nn.LayerNorm(256),
@@ -96,12 +93,12 @@ class ChaoticCoherentGenerator(nn.Module):
             nn.Linear(256, 6)  # [center_x, center_y, scale, rotation, flow_x, flow_y]
         )
 
-        # Projection to spatial maps (separate for each frequency)
+        # projection
         self.low_to_spatial = nn.Linear(512, 8 * 8 * 128)
         self.mid_to_spatial = nn.Linear(512, 16 * 16 * 64)
         self.high_to_spatial = nn.Linear(512, 32 * 32 * 32)
 
-        # Low-frequency path (8x8 -> 32x32)
+        # Low(8x8 -> 32x32)
         self.low_decoder = nn.Sequential(
             nn.ConvTranspose2d(128, 64, 4, 2, 1),  # 8->16
             nn.GroupNorm(8, 64),
@@ -111,14 +108,14 @@ class ChaoticCoherentGenerator(nn.Module):
             nn.GELU()
         )
 
-        # Mid-frequency path (16x16 -> 32x32)
+        # Mid(16x16 -> 32x32)
         self.mid_decoder = nn.Sequential(
             nn.ConvTranspose2d(64, 32, 4, 2, 1),   # 16->32
             nn.GroupNorm(8, 32),
             nn.GELU()
         )
 
-        # High-frequency path (already 32x32)
+        # High pth (32x32)
         self.high_decoder = nn.Sequential(
             nn.Conv2d(32, 32, 3, 1, 1),
             nn.GroupNorm(8, 32),
@@ -126,7 +123,7 @@ class ChaoticCoherentGenerator(nn.Module):
         )
 
         self.fusion = nn.Sequential(
-            nn.Conv2d(96, 64, 3, 1, 1),  # 32+32+32 channels
+            nn.Conv2d(96, 64, 3, 1, 1),
             nn.GroupNorm(8, 64),
             nn.GELU(),
             nn.Conv2d(64, 32, 3, 1, 1),
@@ -152,14 +149,12 @@ class ChaoticCoherentGenerator(nn.Module):
                     nn.init.constant_(m.bias, 0)
 
     def generate_perlin_noise(self, batch_size, size, device):
-        """Generate smooth Perlin-like noise for organic texture"""
         grid_size = 4
         grid = torch.randn(batch_size, 2, grid_size, grid_size, device=device)
         noise = F.interpolate(grid, size=(size, size), mode='bicubic', align_corners=True)
-        return noise[:, 0:1]  # Single channel
+        return noise[:, 0:1]  # single channel
 
     def apply_color_palette(self, grayscale, palette):
-        """Map grayscale values to feature-driven color palette"""
         batch_size = grayscale.size(0)
         palette = palette.view(batch_size, 4, 3)  # 4 clors, 3 channels
 
@@ -172,7 +167,6 @@ class ChaoticCoherentGenerator(nn.Module):
             channel_colors = palette[:, :, i]  # (batch, 4)
             stops = torch.linspace(0, 1, 4, device=grayscale.device)
 
-            # Interpolate between color stops
             result = torch.zeros_like(gray_norm)
             for j in range(3):
                 mask = (gray_norm >= stops[j]) & (gray_norm < stops[j+1])
@@ -191,37 +185,36 @@ class ChaoticCoherentGenerator(nn.Module):
         batch_size = x.size(0)
         device = x.device
 
-        # Extract feature groups
+        # extract feature groups
         structure = x[:, self.structure_idx]
         control = x[:, self.control_idx]
         operations = x[:, self.operations_idx]
         style = x[:, self.style_idx]
         advanced = x[:, self.advanced_idx]
 
-        # Encode to frequency bands
+        # encode to frequency bands
         low_freq = self.low_freq_encoder(structure)
         mid_freq = self.mid_freq_encoder(torch.cat([control, operations], dim=1))
         high_freq = self.high_freq_encoder(torch.cat([style, advanced], dim=1))
 
-        # Generate color palette
+        # generate color palette
         palette = self.color_palette_net(x)
 
-        # Generate composition parameters
+        # generate composition parameters
         composition_params = self.composition_net(low_freq)
         composition_params = torch.tanh(composition_params)
 
-        # Project to spatial maps
+        # Project
         low_spatial = self.low_to_spatial(low_freq).view(batch_size, 128, 8, 8)
         mid_spatial = self.mid_to_spatial(mid_freq).view(batch_size, 64, 16, 16)
         high_spatial = self.high_to_spatial(high_freq).view(batch_size, 32, 32, 32)
 
-        # Decode each frequency band
+        # Decode
         low_decoded = self.low_decoder(low_spatial)     # (batch, 32, 32, 32)
         mid_decoded = self.mid_decoder(mid_spatial)     # (batch, 32, 32, 32)
         high_decoded = self.high_decoder(high_spatial)  # (batch, 32, 32, 32)
 
-        # Apply spatial transformations based on composition params
-        # (rotation, translation via grid_sample)
+        # Apply 
         theta = composition_params[:, 3:4]  # rotation
         scale = 0.5 + composition_params[:, 2:3]  # scale
         tx, ty = composition_params[:, 0:1], composition_params[:, 1:2]
@@ -239,32 +232,42 @@ class ChaoticCoherentGenerator(nn.Module):
         grid = F.affine_grid(affine_matrix, low_decoded.size(), align_corners=False)
         low_decoded = F.grid_sample(low_decoded, grid, align_corners=False, padding_mode='reflection')
 
-        # Merge frequency bands
         merged = torch.cat([low_decoded, mid_decoded, high_decoded], dim=1)
-        fused = self.fusion(merged)  # (batch, 3, 32, 32)
+        fused = self.fusion(merged)
 
-        # Add controlled Perlin noise for organic texture
+        # Add controlled Perlin noise 
         noise_scale = self.noise_strength(x)
         perlin = self.generate_perlin_noise(batch_size, 32, device)
         perlin = perlin.expand(-1, 3, -1, -1)
         fused = fused + noise_scale.view(-1, 1, 1, 1) * perlin * 0.3
 
-        # Convert to grayscale for palette mapping
+        # Convert to grayscale
         grayscale = fused.mean(dim=1, keepdim=True)
 
-        # Apply feature-driven color palette
         colored = self.apply_color_palette(grayscale, palette)
 
-        # Final normalization
         output = torch.tanh(colored)
 
         return output
 
-
-# ===================== Feature-Aware Loss Functions =====================
+def feature_consistency_loss(features, images):
+    """Ensure feature similarity leads to image similarity"""
+    batch_size = features.size(0)
+    if batch_size < 4:
+        return torch.tensor(0.0, device=features.device)
+    
+    feat_dists = torch.cdist(features, features)
+    img_flat = images.view(batch_size, -1)
+    img_dists = torch.cdist(img_flat, img_flat)
+    
+    feat_dists = feat_dists / (feat_dists.max() + 1e-8)
+    img_dists = img_dists / (img_dists.max() + 1e-8)
+    
+    consistency = F.mse_loss(feat_dists, img_dists)
+    
+    return consistency
 
 def frequency_coherence_loss(images):
-    """Enforce low-frequency dominance with high-frequency details"""
     # FFT decomposition
     fft = torch.fft.fft2(images)
     fft_shifted = torch.fft.fftshift(fft)
@@ -290,7 +293,6 @@ def frequency_coherence_loss(images):
 
 
 def color_harmony_loss(images):
-    """Encourage harmonious color distributions"""
     batch_size = images.size(0)
 
     # convert to LAB-like space (approximation)
@@ -308,7 +310,6 @@ def color_harmony_loss(images):
 
 
 def texture_diversity_loss(images):
-    """Encourage multi-scale texture patterns"""
     textures = []
     for scale in [1, 2, 4]:
         if scale > 1:
@@ -324,9 +325,10 @@ def texture_diversity_loss(images):
 
     return F.relu(0.1 - diversity)  # Penalize if diversity < 0.1
 
+def safe_item(x):
+    return x.item() if isinstance(x, torch.Tensor) else float(x)
 
 def combined_aesthetic_loss(features, images):
-    """Enhanced loss with BALANCED weights"""
 
     consistency = feature_consistency_loss(features, images)
     variance = feature_variance_preservation_loss(features, images)
@@ -339,21 +341,21 @@ def combined_aesthetic_loss(features, images):
     smoothness = tv_h + tv_w
 
     total = (
-        10.0 * consistency +      # ~1.5 → 15
-        0.05 * variance +         # ~50 → 2.5
-        0.5 * freq_coherence +    # ~14 → 7
-        0.2 * color_harmony +     # ~31 → 6.2
-        0.01 * texture_div +      # ~126 → 1.26
-        0.1 * smoothness          # ~38 → 3.8
+        10.0 * consistency +      
+        0.05 * variance +         
+        0.5 * freq_coherence +    
+        0.2 * color_harmony +     
+        0.01 * texture_div +     
+        0.1 * smoothness       
     )
 
     return total, {
-        'consistency': consistency.item(),
-        'variance': variance.item(),
-        'freq_coherence': freq_coherence.item(),
-        'color_harmony': color_harmony.item(),
-        'texture_div': texture_div.item(),
-        'smoothness': smoothness.item()
+        'consistency': safe_item(consistency),
+        'variance': safe_item(variance),
+        'freq_coherence': safe_item(freq_coherence),
+        'color_harmony': safe_item(color_harmony),
+        'texture_div': safe_item(texture_div),
+        'smoothness': safe_item(smoothness)
     }
 
 def frequency_coherence_loss(images):
@@ -385,7 +387,7 @@ def frequency_coherence_loss(images):
 
 
 def feature_variance_preservation_loss(features, images):
-    """Ensure each feature dimension contributes to image variance"""
+    """this function ensure each feature dimension contributes to image variance"""
     batch_size = features.size(0)
     if batch_size < 4:
         return torch.tensor(0.0, device=features.device)
@@ -417,22 +419,19 @@ def hsv_to_rgb(h, s, v):
 
 # ===================== Trainer =====================
 def create_feature_explanation_grid(model, features, save_path):
-    """
-    Show how specific features affect the output
-    Systematic perturbation of individual features
-    """
+
     model.eval()
-    base_idx = 42  # Choose interesting sample
+    base_idx = 42  # random sample
     base_features = features[base_idx].copy()
 
-    # Select 6 interesting features to visualize
+    #  6  features to visualize
     feature_indices = [0, 3, 10, 14, 19, 24]  # LOC, loops, nesting, comments, complexity, indentation
     feature_names = [FEATURE_NAMES[i] for i in feature_indices]
 
     fig, axes = plt.subplots(len(feature_indices), 5, figsize=(15, 3*len(feature_indices)))
 
     for row, feat_idx in enumerate(feature_indices):
-        # Create variations: [-2σ, -1σ, 0, +1σ, +2σ]
+        # create variations: [-2σ, -1σ, 0, +1σ, +2σ]
         variations = []
         for multiplier in [-2, -1, 0, 1, 2]:
             varied = base_features.copy()
@@ -466,28 +465,23 @@ def compute_color_palette_from_features(features):
     # Normalize to [0, 1]
     f_norm = (features - features.min(0)) / (features.max(0) - features.min(0) + 1e-8)
 
-    # Base hue from structure
     structure_score = f_norm[[0, 1, 2]].mean()
     base_hue = structure_score * 0.7  # 0 (red) to 0.7 (blue)
 
-    # Saturation from control flow
     control_score = f_norm[[3, 4]].mean()
     saturation = 0.3 + control_score * 0.6
 
-    # Value from style
     style_score = f_norm[[14, 15]].mean()
     value = 0.4 + style_score * 0.5
 
-    # Create 4 colors with variation
     palette = []
     for shift in [-0.1, 0, 0.1, 0.2]:
         hue = (base_hue + shift) % 1.0
         palette.append(hsv_to_rgb(hue, saturation, value))
 
-    return np.array(palette).flatten()  # (12,) vector
+    return np.array(palette).flatten()
 
 def interpolate_code_features(model, features, idx1, idx2, steps=10):
-    """Smooth transition between two code samples"""
     model.eval()
 
     f1 = features[idx1]
@@ -505,7 +499,7 @@ def interpolate_code_features(model, features, idx1, idx2, steps=10):
     return images
 
 
-class FeatureDrivenTrainer:
+class Trainer:
     def __init__(self, model, device='cpu'):
         self.model = model.to(device)
         self.device = device
@@ -564,12 +558,11 @@ class FeatureDrivenTrainer:
             min_lr=1e-6
         )
 
-        print(f"Training chaotic-coherent model for {epochs} epochs on {self.device}...")
-        print("Focus: Frequency coherence + color harmony + texture diversity")
+        print(f"Training model for {epochs} epochs on {self.device}...")
 
         for epoch in range(epochs):
             avg_loss, components = self.train_epoch(train_loader, optimizer)
-            scheduler.step(avg_loss)  # Pass loss for plateau detection
+            scheduler.step(avg_loss) 
 
             if (epoch + 1) % 10 == 0:
                 print(f"Epoch [{epoch+1}/{epochs}] - Loss: {avg_loss:.4f}")
@@ -642,9 +635,7 @@ class FeatureDrivenTrainer:
         print(f"Samples saved to {save_path}")
 
     def test_feature_consistency(self, features, save_path='results/consistency_test.png'):
-        """
-        Visualize that similar features produce similar images.
-        """
+
         os.makedirs(os.path.dirname(save_path), exist_ok=True)
         self.model.eval()
 
@@ -689,7 +680,6 @@ class FeatureDrivenTrainer:
         print(f"Consistency test saved to {save_path}")
 
 
-# ===================== Utilities =====================
 
 def prepare_data_loaders(features, batch_size=32):
     dataset = CodeToImageDataset(features)
@@ -698,19 +688,17 @@ def prepare_data_loaders(features, batch_size=32):
     return loader
 
 
-# ===================== Main =====================
 
 def main():
     os.makedirs("models", exist_ok=True)
     os.makedirs("results", exist_ok=True)
 
     print("=" * 70)
-    print("Feature-Driven Code-to-Image Generator (Final Version)")
-    print("Features → Meaningful Visual Attributes")
+    print(" Code-to-Image Generator (Final Version)")
     print("=" * 70)
 
     # Load features
-    print("\n[1/5] Loading features...")
+    print("\n Loading features...")
     features_path = "data/features.npy"
     if not os.path.exists(features_path):
         raise FileNotFoundError(f"{features_path} not found")
@@ -719,33 +707,31 @@ def main():
     print(f"Features shape: {features.shape}")
 
     # Prepare data
-    print("\n[2/5] Preparing data loaders...")
+    print("\nPreparing data loaders...")
     loader = prepare_data_loaders(features, batch_size=32)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Device: {device}")
 
     # Initialize model
-    print("\n[3/5] Initializing feature-driven CNN model...")
-    model = ChaoticCoherentGenerator(input_dim=features.shape[1], image_size=32)
+    print("\n InitializingCNN model...")
+    model = Generator(input_dim=features.shape[1], image_size=32)
     param_count = sum(p.numel() for p in model.parameters())
     print(f"Model parameters: {param_count:,}")
 
     # Train model
-    print("\n[4/5] Training model...")
-    trainer = FeatureDrivenTrainer(model, device=device)
+    print("\nTraining model...")
+    trainer = Trainer(model, device=device)
     trainer.train(loader, epochs=100, lr=1e-3)
 
     # Save results
     print("\n[5/5] Saving results...")
-    trainer.save_model("models/chotic.pth")
-    trainer.plot_losses("results/chotic_losses.png")
-    trainer.generate_samples(features, n=25, save_path="results/chotic_samples.png")
-    trainer.test_feature_consistency(features, save_path="results/chotic_consistency.png")
+    trainer.save_model("models/model_32.pth")
+    trainer.plot_losses("results/losses.png")
+    trainer.generate_samples(features, n=25, save_path="results/samples.png")
+    trainer.test_feature_consistency(features, save_path="results/consistency.png")
 
     print("\n" + "=" * 70)
-    print("Complete! Check results/ folder.")
-    print("=" * 70)
-
+    print("Check results/ folder.")
 
 if __name__ == "__main__":
     main()