Restore colorspace detection using loaded model

Detection now uses the same model instance that's already in memory
with offload strategy configured. The detection respects the model's
memory management setup.
parent 3bce5085
......@@ -1354,8 +1354,11 @@ def detect_model_colorspace(pipe, model_name, m_info, args):
def get_model_colorspace(pipe, model_name, m_info, args):
"""Get the colorspace for a model, detecting it if necessary.
Uses the already-loaded model to generate a minimal test frame.
The detection respects the model's offload strategy.
Args:
pipe: The loaded diffusion pipeline
pipe: The loaded diffusion pipeline (already configured with offload)
model_name: Name of the model in MODELS config
m_info: Model info dict from MODELS
args: Command line arguments
......@@ -1367,29 +1370,8 @@ def get_model_colorspace(pipe, model_name, m_info, args):
if model_name in MODELS and "colorspace" in MODELS[model_name]:
return MODELS[model_name]["colorspace"]
# Skip detection if using offload strategies (memory-constrained)
# Colorspace detection requires extra memory that may not be available
offload_strategy = getattr(args, 'offload_strategy', 'model')
if offload_strategy not in ['none', 'model']:
print(f" 📊 Colorspace: Skipping detection (offload_strategy={offload_strategy})")
print(f" Defaulting to RGB (use --swap_bgr if colors look wrong)")
# Save RGB as default to avoid repeated attempts
if model_name in MODELS:
MODELS[model_name]["colorspace"] = "RGB"
save_models_config(MODELS)
return "RGB"
# Also skip if low system RAM
available_ram_gb = get_available_ram_gb()
if available_ram_gb < 12.0: # Need ample RAM for detection
print(f" 📊 Colorspace: Skipping detection (low RAM: {available_ram_gb:.1f}GB)")
print(f" Defaulting to RGB (use --swap_bgr if colors look wrong)")
if model_name in MODELS:
MODELS[model_name]["colorspace"] = "RGB"
save_models_config(MODELS)
return "RGB"
# Detect and save
# Run detection using the same model that's already loaded
# This respects the offload strategy already configured
return detect_model_colorspace(pipe, model_name, m_info, args)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment