flux_app/config.py

# config.py

import torch

DTYPE = torch.bfloat16
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BASE_MODEL = "black-forest-labs/FLUX.1-dev"
TAEF1_MODEL = "madebyollin/taef1"
MAX_SEED = 2**32 - 1

# You might want to make these configurable via environment variables:
# import os
# BASE_MODEL = os.environ.get("BASE_MODEL", "black-forest-labs/FLUX.1-dev")
# TAEF1_MODEL = os.environ.get("TAEF1_MODEL", "madebyollin/taef1")
# DEVICE = os.environ.get("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
# DTYPE = os.environ.get("DTYPE", "bfloat16")  # Or "float32", "float16"
# if DTYPE == "bfloat16":
#     DTYPE = torch.bfloat16
# elif DTYPE == "float16":
#     DTYPE = torch.float16
# else:
#    DTYPE = torch.float32