monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move num_train_steps to ModelConfig

Browse Source
This commit is contained in:
filipstrand 2024-08-20 11:45:36 +02:00
parent 7795d47cd3
commit 4d613c7390
2 changed files with 15 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
src/flux_1_schnell/config/model_config.py
View File

@ -2,12 +2,19 @@ from enum import Enum
class ModelConfig(Enum): class ModelConfig(Enum):
FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 512) FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 1000, 512)
FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 256) FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256)
def __init__(self, model_name: str, alias: str, max_sequence_length: int): def __init__(
self,
model_name: str,
alias: str,
num_train_steps: int,
max_sequence_length: int,
):
self.alias = alias self.alias = alias
self.model_name = model_name self.model_name = model_name
self.num_train_steps = num_train_steps
self.max_sequence_length = max_sequence_length self.max_sequence_length = max_sequence_length
@staticmethod @staticmethod

6
src/flux_1_schnell/config/runtime_config.py
View File

@ -9,7 +9,7 @@ class RuntimeConfig:
def __init__(self, config: Config, model_config: ModelConfig): def __init__(self, config: Config, model_config: ModelConfig):
self.config = config self.config = config
self.num_train_steps = 1000 self.model_config = model_config
self.sigmas = self._create_sigmas(config, model_config) self.sigmas = self._create_sigmas(config, model_config)
@property @property
@ -32,6 +32,10 @@ class RuntimeConfig:
def precision(self): def precision(self):
return self.config.precision return self.config.precision
@property
def num_train_steps(self):
return self.model_config.num_train_steps
@staticmethod @staticmethod
def _create_sigmas(config, model): def _create_sigmas(config, model):
sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps) sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)