config.py 1.69 KB
Newer Older
litzh's avatar
litzh committed
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
import os
from dataclasses import dataclass
from pathlib import Path

from loguru import logger


@dataclass
class ServerConfig:
    host: str = "0.0.0.0"
    port: int = 8000
    max_queue_size: int = 10

    task_timeout: int = 300
    task_history_limit: int = 1000

    http_timeout: int = 30
    http_max_retries: int = 3

    cache_dir: str = str(Path(__file__).parent.parent / "server_cache")
    max_upload_size: int = 500 * 1024 * 1024  # 500MB

    lora_dir: str = None  # Directory for LoRA files

    @classmethod
    def from_env(cls) -> "ServerConfig":
        config = cls()

        if env_host := os.environ.get("LIGHTX2V_HOST"):
            config.host = env_host

        if env_port := os.environ.get("LIGHTX2V_PORT"):
            try:
                config.port = int(env_port)
            except ValueError:
                logger.warning(f"Invalid port in environment: {env_port}")

        if env_queue_size := os.environ.get("LIGHTX2V_MAX_QUEUE_SIZE"):
            try:
                config.max_queue_size = int(env_queue_size)
            except ValueError:
                logger.warning(f"Invalid max queue size: {env_queue_size}")

        # MASTER_ADDR is now managed by torchrun, no need to set manually

        if env_cache_dir := os.environ.get("LIGHTX2V_CACHE_DIR"):
            config.cache_dir = env_cache_dir

        return config

    def validate(self) -> bool:
        valid = True

        if self.max_queue_size <= 0:
            logger.error("max_queue_size must be positive")
            valid = False

        if self.task_timeout <= 0:
            logger.error("task_timeout must be positive")
            valid = False

        return valid


server_config = ServerConfig.from_env()