File size: 5,422 Bytes
cc5958e
 
 
 
 
 
 
6287e96
 
 
 
 
 
 
cc5958e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6287e96
fb609fe
 
 
 
cc5958e
 
 
 
 
 
 
fb609fe
 
cc5958e
fb609fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70a95c9
fb609fe
 
 
 
 
 
6287e96
 
fb609fe
6287e96
 
 
 
 
 
 
 
 
 
 
 
 
cc5958e
6287e96
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
"""Configuration management for the Pepe meme generator.

This module defines all configuration parameters for model selection,
generation settings, and application behavior. The ModelConfig dataclass
provides a centralized configuration system with sensible defaults.

"""

from dataclasses import dataclass
from typing import Optional


@dataclass
class ModelConfig:
    """
    Central configuration for model and generation parameters.
    
    This dataclass contains all settings for model selection, generation
    parameters, and optimization flags. It supports multiple models including
    fine-tuned LoRA variants and fast LCM models.
    
    Attributes:
        AVAILABLE_MODELS: Dictionary of available model configurations
        SELECTED_MODEL: Currently selected model name
        BASE_MODEL: HuggingFace ID of the base Stable Diffusion model
        LORA_PATH: Path or HuggingFace ID of LoRA weights
        USE_LORA: Whether to load and use LoRA weights
        USE_LCM: Whether to use LCM (Latent Consistency Model) for fast inference
        LCM_LORA_PATH: Path to LCM-LoRA weights
        TRIGGER_WORD: Trigger word to activate fine-tuned style
        DEFAULT_STEPS: Default number of diffusion steps
        DEFAULT_GUIDANCE: Default guidance scale (CFG)
        DEFAULT_WIDTH: Default output image width
        DEFAULT_HEIGHT: Default output image height
        DEFAULT_NEGATIVE_PROMPT: Default negative prompt for all generations
        FORCE_CPU: Force CPU mode (disable GPU)
        ENABLE_XFORMERS: Enable memory-efficient attention
    """
    
    # Available models
    AVAILABLE_MODELS: dict = None
    
    def __post_init__(self):
        """
        Initialize AVAILABLE_MODELS dictionary if not already set.
        
        This method is called automatically after __init__. It populates
        the AVAILABLE_MODELS dictionary with all supported model configurations.
        Each model can have different base models, LoRA weights, and optimization flags.
        """
        if self.AVAILABLE_MODELS is None:
            self.AVAILABLE_MODELS = {
                # Primary fine-tuned model - Best quality, trained on Pepe dataset
                "Pepe Fine-tuned (LoRA)": {
                    "base": "runwayml/stable-diffusion-v1-5",
                    "lora": "MJaheen/Pepe_The_Frog_model_v1_lora",
                    "trigger_word": "pepe_style_frog",
                    "use_lora": True,
                    "use_lcm": False
                },
                "Pepe + LCM (FAST)": {
                    "base": "runwayml/stable-diffusion-v1-5",
                    "lora": "MJaheen/Pepe_The_Frog_model_v1_lora",
                    "lcm_lora": "latent-consistency/lcm-lora-sdv1-5",
                    "trigger_word": "pepe_style_frog",
                    "use_lora": True,
                    "use_lcm": True
                },
                "Base SD 1.5": {
                    "base": "runwayml/stable-diffusion-v1-5",
                    "lora": None,
                    "trigger_word": "pepe the frog",
                    "use_lora": False,
                    "use_lcm": False
                },
                "Dreamlike Photoreal 2.0": {
                    "base": "dreamlike-art/dreamlike-photoreal-2.0",
                    "lora": None,
                    "trigger_word": "pepe the frog",
                    "use_lora": False,
                    "use_lcm": False
                },
                "Openjourney v4": {
                    "base": "prompthero/openjourney-v4",
                    "lora": None,
                    "trigger_word": "pepe the frog",
                    "use_lora": False,
                    "use_lcm": False
                },
                "Tiny SD (Fast CPU)": {
                    "base": "segmind/tiny-sd",
                    "lora": None,
                    "trigger_word": "pepe the frog",
                    "use_lora": False,
                    "use_lcm": False
                },
                "Small SD (Balanced CPU)": {
                    "base": "segmind/small-sd",
                    "lora": None,
                    "trigger_word": "pepe the frog",
                    "use_lora": False,
                    "use_lcm": False
                }
            }
    
    # Default model selection
    SELECTED_MODEL: str = "Pepe Fine-tuned (LoRA)"
    
    # Model paths (will be set based on selection)
    BASE_MODEL: str = "runwayml/stable-diffusion-v1-5"
    LORA_PATH: str = "MJaheen/Pepe_The_Frog_model_v1_lora"
    USE_LORA: bool = True
    TRIGGER_WORD: str = "pepe_style_frog"
    
    # LCM settings
    USE_LCM: bool = False
    LCM_LORA_PATH: Optional[str] = None
    
    # Default generation parameters
    DEFAULT_STEPS: int = 25  # Reduced for faster CPU inference (was 50)
    DEFAULT_GUIDANCE: float = 7.5
    DEFAULT_WIDTH: int = 512
    DEFAULT_HEIGHT: int = 512
    
    # Negative prompt
    DEFAULT_NEGATIVE_PROMPT: str = (
        "blurry, low quality, distorted, deformed, "
        "ugly, bad anatomy, watermark, signature"
    )
    
    # Performance
    ENABLE_ATTENTION_SLICING: bool = True
    ENABLE_VAE_SLICING: bool = True
    FORCE_CPU: bool = False  # Set to True to force CPU, False to use GPU if available
    
    # Available styles
    AVAILABLE_STYLES: tuple = (
        "default", "happy", "sad", "smug", 
        "angry", "thinking", "surprised"
    )