config.py

"""
Configuration file for Interior Style Transfer Pipeline
"""
import os
from pathlib import Path

class Config:
    # Model paths and settings
    SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
    SDXL_REFINER_ID = "stabilityai/stable-diffusion-xl-refiner-1.0"
    CONTROLNET_MODEL_ID = "lllyasviel/control_v11p_sd15_seg"
    
    # Generation parameters
    NUM_INFERENCE_STEPS = 50
    GUIDANCE_SCALE = 7.5
    STRENGTH = 0.8  # How much to preserve original structure
    
    # Image processing
    IMAGE_SIZE = 1024
    BATCH_SIZE = 1
    
    # Segmentation settings
    SEGMENTATION_CONFIDENCE = 0.5
    PRESERVE_CLASSES = ['wall', 'window', 'door', 'ceiling', 'floor']
    
    # Style transfer parameters
    STYLE_STRENGTH = 0.7
    CONTENT_STRENGTH = 0.3
    BLEND_ALPHA = 0.8
    
    # Output settings
    OUTPUT_DIR = "outputs"
    TEMP_DIR = "temp"
    
    # Device
    DEVICE = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
    
    @classmethod
    def create_directories(cls):
        """Create necessary directories"""
        Path(cls.OUTPUT_DIR).mkdir(exist_ok=True)
        Path(cls.TEMP_DIR).mkdir(exist_ok=True)