John Pollock
6cb71ac51c
This commit introduces block swapping functionality for Qwen models, enabling them to run on systems with limited VRAM by offloading layers to a swap device (e.g., CPU RAM). Key changes: - A new `QwenBlockSwapManager` class is implemented to handle the patching of Qwen transformer blocks. - The `apply_block_swap` function is extended to detect Qwen models and apply the swapping logic to their `transformer_blocks`. - A model signature for Qwen is added to `model_sig.py` to correctly identify the swappable modules. - A new diagnostic function, `log_unsupported_model_analysis`, is added to log the structure of unsupported models, aiding future development.
22 lines
864 B
Python
22 lines
864 B
Python
def get_model_type(model_patcher):
|
|
"""
|
|
Identifies the model type using a multi-layered approach for robustness.
|
|
It first checks the diffusion model's class name, then falls back to the
|
|
model_type enum.
|
|
"""
|
|
if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'diffusion_model'):
|
|
class_name = type(model_patcher.model.diffusion_model).__name__
|
|
|
|
# Prioritize class name for accuracy
|
|
if "Flux" in class_name:
|
|
return "FLUX"
|
|
if "Qwen" in class_name:
|
|
return "QWEN"
|
|
|
|
# Fallback to the model_type enum for other cases
|
|
if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'model_type'):
|
|
# model_type is an enum, so we return its name as a string
|
|
return model_patcher.model.model_type.name
|
|
|
|
return "UNKNOWN"
|