haagenti_sparse/
categories.rs

1//! Attention head and prompt category definitions
2
3use serde::{Deserialize, Serialize};
4use smallvec::SmallVec;
5use std::collections::HashMap;
6
7/// Categories of attention heads based on their learned function
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub enum HeadCategory {
10    /// Face/portrait focused heads (eyes, mouth, structure)
11    Face,
12    /// Body and pose attention
13    Body,
14    /// Background and scene composition
15    Background,
16    /// Style and texture patterns
17    Style,
18    /// Composition and layout
19    Composition,
20    /// Fine detail refinement
21    Detail,
22    /// Lighting and shadows
23    Lighting,
24    /// Color harmony
25    Color,
26    /// Object boundaries
27    Edge,
28    /// General purpose (always active)
29    General,
30}
31
32impl HeadCategory {
33    /// Default importance for this category (0.0 - 1.0)
34    pub fn default_importance(&self) -> f32 {
35        match self {
36            HeadCategory::General => 1.0,
37            HeadCategory::Composition => 0.9,
38            HeadCategory::Style => 0.8,
39            HeadCategory::Detail => 0.7,
40            HeadCategory::Face => 0.6,
41            HeadCategory::Body => 0.6,
42            HeadCategory::Background => 0.5,
43            HeadCategory::Lighting => 0.5,
44            HeadCategory::Color => 0.4,
45            HeadCategory::Edge => 0.4,
46        }
47    }
48
49    /// Whether this category should always be active
50    pub fn is_mandatory(&self) -> bool {
51        matches!(self, HeadCategory::General | HeadCategory::Composition)
52    }
53
54    /// All categories
55    pub fn all() -> &'static [HeadCategory] {
56        &[
57            HeadCategory::Face,
58            HeadCategory::Body,
59            HeadCategory::Background,
60            HeadCategory::Style,
61            HeadCategory::Composition,
62            HeadCategory::Detail,
63            HeadCategory::Lighting,
64            HeadCategory::Color,
65            HeadCategory::Edge,
66            HeadCategory::General,
67        ]
68    }
69}
70
71/// Categories of prompts that influence head activation
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
73pub enum PromptCategory {
74    /// Portrait or character-focused
75    Portrait,
76    /// Landscape or scene
77    Landscape,
78    /// Abstract or artistic
79    Abstract,
80    /// Photorealistic
81    Photorealistic,
82    /// Anime/cartoon style
83    Anime,
84    /// Architecture or interiors
85    Architecture,
86    /// Objects or products
87    Object,
88    /// Fantasy or sci-fi
89    Fantasy,
90    /// Unknown/mixed
91    Mixed,
92}
93
94impl PromptCategory {
95    /// Keywords that indicate this category
96    pub fn keywords(&self) -> &'static [&'static str] {
97        match self {
98            PromptCategory::Portrait => &[
99                "portrait",
100                "face",
101                "person",
102                "woman",
103                "man",
104                "girl",
105                "boy",
106                "character",
107                "headshot",
108                "bust",
109                "selfie",
110            ],
111            PromptCategory::Landscape => &[
112                "landscape",
113                "mountain",
114                "forest",
115                "ocean",
116                "sky",
117                "sunset",
118                "sunrise",
119                "nature",
120                "scenery",
121                "vista",
122                "horizon",
123            ],
124            PromptCategory::Abstract => &[
125                "abstract",
126                "geometric",
127                "pattern",
128                "fractal",
129                "surreal",
130                "dreamlike",
131                "psychedelic",
132                "minimalist",
133            ],
134            PromptCategory::Photorealistic => &[
135                "photo",
136                "photograph",
137                "realistic",
138                "hyperrealistic",
139                "photorealistic",
140                "raw",
141                "unedited",
142                "natural",
143            ],
144            PromptCategory::Anime => &[
145                "anime",
146                "manga",
147                "cartoon",
148                "illustrated",
149                "cel-shaded",
150                "2d",
151                "chibi",
152                "kawaii",
153            ],
154            PromptCategory::Architecture => &[
155                "building",
156                "architecture",
157                "interior",
158                "room",
159                "house",
160                "skyscraper",
161                "cathedral",
162                "bridge",
163                "structure",
164            ],
165            PromptCategory::Object => &[
166                "product",
167                "object",
168                "item",
169                "thing",
170                "device",
171                "tool",
172                "furniture",
173                "vehicle",
174                "food",
175            ],
176            PromptCategory::Fantasy => &[
177                "fantasy",
178                "magical",
179                "dragon",
180                "wizard",
181                "elf",
182                "fairy",
183                "mythical",
184                "enchanted",
185                "sci-fi",
186                "futuristic",
187                "cyberpunk",
188            ],
189            PromptCategory::Mixed => &[],
190        }
191    }
192
193    /// Head category importance modifiers for this prompt category
194    pub fn category_weights(&self) -> HashMap<HeadCategory, f32> {
195        let mut weights = HashMap::new();
196
197        match self {
198            PromptCategory::Portrait => {
199                weights.insert(HeadCategory::Face, 1.0);
200                weights.insert(HeadCategory::Body, 0.8);
201                weights.insert(HeadCategory::Background, 0.3);
202                weights.insert(HeadCategory::Style, 0.7);
203                weights.insert(HeadCategory::Detail, 0.9);
204                weights.insert(HeadCategory::Lighting, 0.7);
205            }
206            PromptCategory::Landscape => {
207                weights.insert(HeadCategory::Face, 0.1);
208                weights.insert(HeadCategory::Body, 0.2);
209                weights.insert(HeadCategory::Background, 1.0);
210                weights.insert(HeadCategory::Composition, 1.0);
211                weights.insert(HeadCategory::Lighting, 0.9);
212                weights.insert(HeadCategory::Color, 0.8);
213            }
214            PromptCategory::Abstract => {
215                weights.insert(HeadCategory::Face, 0.0);
216                weights.insert(HeadCategory::Body, 0.0);
217                weights.insert(HeadCategory::Style, 1.0);
218                weights.insert(HeadCategory::Composition, 1.0);
219                weights.insert(HeadCategory::Color, 0.9);
220                weights.insert(HeadCategory::Edge, 0.7);
221            }
222            PromptCategory::Photorealistic => {
223                weights.insert(HeadCategory::Detail, 1.0);
224                weights.insert(HeadCategory::Lighting, 1.0);
225                weights.insert(HeadCategory::Color, 0.9);
226                weights.insert(HeadCategory::Edge, 0.8);
227            }
228            PromptCategory::Anime => {
229                weights.insert(HeadCategory::Face, 0.9);
230                weights.insert(HeadCategory::Style, 1.0);
231                weights.insert(HeadCategory::Edge, 0.9);
232                weights.insert(HeadCategory::Color, 0.8);
233                weights.insert(HeadCategory::Detail, 0.5);
234            }
235            PromptCategory::Architecture => {
236                weights.insert(HeadCategory::Face, 0.0);
237                weights.insert(HeadCategory::Composition, 1.0);
238                weights.insert(HeadCategory::Edge, 1.0);
239                weights.insert(HeadCategory::Detail, 0.9);
240                weights.insert(HeadCategory::Lighting, 0.8);
241            }
242            PromptCategory::Object => {
243                weights.insert(HeadCategory::Face, 0.0);
244                weights.insert(HeadCategory::Detail, 1.0);
245                weights.insert(HeadCategory::Edge, 0.9);
246                weights.insert(HeadCategory::Lighting, 0.8);
247                weights.insert(HeadCategory::Color, 0.7);
248            }
249            PromptCategory::Fantasy => {
250                weights.insert(HeadCategory::Style, 1.0);
251                weights.insert(HeadCategory::Composition, 0.9);
252                weights.insert(HeadCategory::Lighting, 0.9);
253                weights.insert(HeadCategory::Color, 0.8);
254            }
255            PromptCategory::Mixed => {
256                // Use default weights
257            }
258        }
259
260        weights
261    }
262
263    /// Detect category from prompt text
264    pub fn detect(prompt: &str) -> SmallVec<[PromptCategory; 3]> {
265        let prompt_lower = prompt.to_lowercase();
266        let mut scores: Vec<(PromptCategory, usize)> = Vec::new();
267
268        for category in &[
269            PromptCategory::Portrait,
270            PromptCategory::Landscape,
271            PromptCategory::Abstract,
272            PromptCategory::Photorealistic,
273            PromptCategory::Anime,
274            PromptCategory::Architecture,
275            PromptCategory::Object,
276            PromptCategory::Fantasy,
277        ] {
278            let count = category
279                .keywords()
280                .iter()
281                .filter(|kw| prompt_lower.contains(*kw))
282                .count();
283
284            if count > 0 {
285                scores.push((*category, count));
286            }
287        }
288
289        scores.sort_by(|a, b| b.1.cmp(&a.1));
290
291        let mut result: SmallVec<[PromptCategory; 3]> =
292            scores.into_iter().take(3).map(|(cat, _)| cat).collect();
293
294        if result.is_empty() {
295            result.push(PromptCategory::Mixed);
296        }
297
298        result
299    }
300}
301
302/// Mapping from head indices to categories
303#[derive(Debug, Clone, Serialize, Deserialize)]
304pub struct CategoryMapping {
305    /// Number of heads per layer
306    pub num_heads: usize,
307    /// Number of layers
308    pub num_layers: usize,
309    /// Category for each head in each layer
310    pub head_categories: Vec<Vec<HeadCategory>>,
311    /// Model this mapping was trained for
312    pub model_id: String,
313}
314
315impl CategoryMapping {
316    /// Create default mapping for SDXL-like architecture
317    pub fn sdxl_default() -> Self {
318        let num_heads = 32;
319        let num_layers = 70;
320
321        // Default category assignment based on typical SDXL behavior
322        let head_categories: Vec<Vec<HeadCategory>> = (0..num_layers)
323            .map(|layer| {
324                (0..num_heads)
325                    .map(|head| Self::default_head_category(layer, head, num_layers, num_heads))
326                    .collect()
327            })
328            .collect();
329
330        Self {
331            num_heads,
332            num_layers,
333            head_categories,
334            model_id: "sdxl-base".into(),
335        }
336    }
337
338    /// Default category based on layer and head position
339    fn default_head_category(
340        layer: usize,
341        head: usize,
342        num_layers: usize,
343        num_heads: usize,
344    ) -> HeadCategory {
345        let layer_fraction = layer as f32 / num_layers as f32;
346        let head_fraction = head as f32 / num_heads as f32;
347
348        // Early layers: composition and structure
349        if layer_fraction < 0.2 {
350            return match head % 4 {
351                0 => HeadCategory::Composition,
352                1 => HeadCategory::General,
353                2 => HeadCategory::Edge,
354                _ => HeadCategory::Style,
355            };
356        }
357
358        // Middle layers: content-specific
359        if layer_fraction < 0.7 {
360            if head_fraction < 0.3 {
361                return HeadCategory::Face;
362            } else if head_fraction < 0.5 {
363                return HeadCategory::Body;
364            } else if head_fraction < 0.7 {
365                return HeadCategory::Background;
366            } else {
367                return HeadCategory::Lighting;
368            }
369        }
370
371        // Late layers: detail and refinement
372        match head % 5 {
373            0 => HeadCategory::Detail,
374            1 => HeadCategory::Color,
375            2 => HeadCategory::Edge,
376            3 => HeadCategory::Style,
377            _ => HeadCategory::General,
378        }
379    }
380
381    /// Get category for a specific head
382    pub fn get_category(&self, layer: usize, head: usize) -> Option<HeadCategory> {
383        self.head_categories
384            .get(layer)
385            .and_then(|heads| heads.get(head))
386            .copied()
387    }
388
389    /// Get all heads in a category for a layer
390    pub fn heads_in_category(&self, layer: usize, category: HeadCategory) -> Vec<usize> {
391        self.head_categories
392            .get(layer)
393            .map(|heads| {
394                heads
395                    .iter()
396                    .enumerate()
397                    .filter(|(_, c)| **c == category)
398                    .map(|(i, _)| i)
399                    .collect()
400            })
401            .unwrap_or_default()
402    }
403}
404
405#[cfg(test)]
406mod tests {
407    use super::*;
408
409    #[test]
410    fn test_prompt_category_detection() {
411        let portrait = PromptCategory::detect("A beautiful portrait of a young woman");
412        assert_eq!(portrait[0], PromptCategory::Portrait);
413
414        let landscape = PromptCategory::detect("Mountain landscape at sunset");
415        assert_eq!(landscape[0], PromptCategory::Landscape);
416
417        let anime = PromptCategory::detect("Anime girl with blue hair");
418        assert!(anime.contains(&PromptCategory::Anime));
419
420        // "xyz abc" has no keyword matches, so should return Mixed
421        let mixed = PromptCategory::detect("xyz abc");
422        assert_eq!(mixed[0], PromptCategory::Mixed);
423    }
424
425    #[test]
426    fn test_category_weights() {
427        let portrait_weights = PromptCategory::Portrait.category_weights();
428        assert_eq!(portrait_weights.get(&HeadCategory::Face), Some(&1.0));
429        assert!(portrait_weights.get(&HeadCategory::Background).unwrap() < &0.5);
430
431        let landscape_weights = PromptCategory::Landscape.category_weights();
432        assert_eq!(landscape_weights.get(&HeadCategory::Background), Some(&1.0));
433        assert!(landscape_weights.get(&HeadCategory::Face).unwrap() < &0.2);
434    }
435
436    #[test]
437    fn test_sdxl_mapping() {
438        let mapping = CategoryMapping::sdxl_default();
439        assert_eq!(mapping.num_heads, 32);
440        assert_eq!(mapping.num_layers, 70);
441
442        // Should have a category for every head
443        for layer in &mapping.head_categories {
444            assert_eq!(layer.len(), 32);
445        }
446    }
447}
haagenti_sparse/categories.rs

haagenti_sparse/
categories.rs
