File size: 3,358 Bytes
40440d6 | 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 | """
Configuration system for enhanced RIFT model.
"""
from dataclasses import dataclass, field
from typing import Literal
@dataclass
class EncoderConfig:
type: Literal["siglip", "dinov2", "evaclip"] = "dinov2"
model_size: Literal["small", "base", "large"] = "base"
output_dim: int = 512
depth_aware: bool = False
@dataclass
class LossConfig:
"""
Configuration for loss functions and weights.
"""
recon_weight: float = 10.0 # Charbonnier Loss
color_match_weight: float = 5.0 # ColorMatchingLoss
day_preserve_weight: float = 3.0 # DayPreservationLoss
moe_weight: float = 0.05 # MoE Load Balancing
brightness_cont_weight: float = 0.1
color_cont_weight: float = 0.1
use_charbonnier: bool = True
use_vgg: bool = True
vgg_weight: float = 0.5
# Tuning guide:
# 0.1 → subtle sharpening, safe
# 0.5 → noticeable edge/texture improvement ← default
use_lpips: bool = False
lpips_weight: float = 0.0
use_l1: bool = False
use_gan: bool = False
gan_weight: float = 0.0
# Legacy (disabled)
use_smooth_transition: bool = False
smooth_weight: float = 0.0
use_content_preservation: bool = False
content_weight: float = 0.0
use_brightness_preservation: bool = False
brightness_weight: float = 0.0
min_brightness: float = 0.0
max_brightness: float = 0.0
charbonnier_eps: float = 1e-4
@dataclass
class AttentionConfig:
use_cross_attention: bool = True
use_adain: bool = True
use_depth_attention: bool = True
use_multi_scale_attention: bool = False
@dataclass
class MoEConfig:
enabled: bool = True
num_experts: int = 4
top_k: int = 2
temperature: float = 0.3
noise_std: float = 0.1
load_balance_coef: float = 0.01
@dataclass
class SamplingConfig:
method: Literal["euler", "heun", "consistency"] = "heun"
num_steps: int = 30
use_flow_matching: bool = True
use_consistency: bool = False
use_cfg: bool = True
cfg_scale: float = 7.5
sigma_min: float = 0.002
sigma_max: float = 80.0
@dataclass
class ModelConfig:
in_channels: int = 16
hidden_size: int = 1024
num_double_layers: int = 8
num_single_layers: int = 16
num_heads: int = 16
style_dim: int = 512
use_text: bool = False
text_dim: int = 1152
encoder: EncoderConfig = field(default_factory=EncoderConfig)
loss: LossConfig = field(default_factory=LossConfig)
attention: AttentionConfig = field(default_factory=AttentionConfig)
moe: MoEConfig = field(default_factory=MoEConfig)
sampling: SamplingConfig = field(default_factory=SamplingConfig)
@dataclass
class TrainingConfig:
content_dir: str = "/workspace/ChronFLUX/rafael_ds_ver2/train/day"
style_dir: str = "/workspace/ChronFLUX/rafael_ds_ver2/train/night"
image_size: int = 384
batch_size: int = 1
grad_accum_steps: int = 4
epochs: int = 43
lr_max: float = 1e-7
lr_min: float = 1e-5
weight_decay: float = 1e-4
save_interval_epochs: int = 1
device: str = "cuda"
dtype: str = "bfloat16"
def get_q1_paper_config() -> ModelConfig:
return ModelConfig()
PRESETS = {
"q1_paper": get_q1_paper_config,
}
def get_config(preset: str = "q1_paper") -> ModelConfig:
return PRESETS.get(preset, get_q1_paper_config)()
|