Fixed reForge for Sage-Attention-2++

"use_sage_attention": OptionInfo(False, "Use SageAttention for cross attention layers (SDXL only)").info("Enable SageAttention optimization for faster generation. Only works with SDXL models.").needs_reload_ui(),

"quicksettings_list": OptionInfo(["sd_model_checkpoint", "sd_vae", "CLIP_stop_at_last_layers", "use_sage_attention"], ...

# SageAttention設定
from ldm_patched.modules.args_parser import args

def get_use_sage_attention():
    """動的にSageAttention設定を取得"""
    cmd_line_sage = getattr(args, 'use_sage_attention', False)
    try:
        from modules.shared_cmd_options import cmd_opts
        ui_sage = getattr(cmd_opts, 'use_sage_attention', False)
        return cmd_line_sage or ui_sage
    except:
        return cmd_line_sage

# 後方互換性のため
USE_SAGE_ATTENTION = get_use_sage_attention()

if "use_sage_attention" not in quicksettings_list:
    quicksettings_list.append("use_sage_attention")

def sync_sageattention_settings():
    """shared.optsの変更をcmd_optsに同期し、ログフラグをリセット"""
    sage_enabled = getattr(shared.opts, 'use_sage_attention', False)
    cmd_opts.use_sage_attention = sage_enabled
    
    # ログフラグをリセット
    from ldm_patched.ldm.modules import attention
    if hasattr(attention, 'SAGE_LOGGED_THIS_GEN'):
        attention.SAGE_LOGGED_THIS_GEN = False

original_opts_set = shared.opts.set
def hooked_opts_set(key, value):
    result = original_opts_set(key, value)
    if key == 'use_sage_attention':
        sync_sageattention_settings()

if "use_sage_attention" not in quicksettings_list:
    quicksettings_list.append("use_sage_attention")

def sync_sageattention_settings():
    """shared.optsの変更をcmd_optsに同期し、ログフラグをリセット"""
    sage_enabled = getattr(shared.opts, 'use_sage_attention', False)
    cmd_opts.use_sage_attention = sage_enabled
    
    # ログフラグをリセット
    from ldm_patched.ldm.modules import attention
    if hasattr(attention, 'SAGE_LOGGED_THIS_GEN'):
        attention.SAGE_LOGGED_THIS_GEN = False

original_opts_set = shared.opts.set
def hooked_opts_set(key, value):
    result = original_opts_set(key, value)
    if key == 'use_sage_attention':
        sync_sageattention_settings()

attn_group.add_argument("--use-sage-attention", action="store_true", help="Use sage attention.")

class SdOptimizationSageAttention(SdOptimization):
    name = "sageattention"
    label = "SageAttention"
    cmd_opt = "use_sage_attention"
    priority = 110

    def is_available(self):
        # SageAttention is only available for SDXL models
        if not shared.cmd_opts.use_sage_attention or not shared.sageattention_available or not torch.cuda.is_available():
            return False
        
        # Check if current model is SDXL
        if hasattr(shared, 'sd_model') and shared.sd_model is not None:
            if hasattr(shared.sd_model, 'is_sdxl') and shared.sd_model.is_sdxl:
                return True
            else:
                print("[SageAttention] SageAttention is only supported for SDXL models. Current model is not SDXL.")
                return False
        
        return True

def sageattention_attention_forward(self, x, context=None, mask=None, **kwargs):
    # SageAttention専用のforward実装
    try:
        import sageattention
        print("[SageAttention] SageAttention 2.2 enabled (HND layout)")
        # ... SageAttention実装 ...
    except Exception as e:
        # フォールバック処理
        
def sageattention_attnblock_forward(self, x):
    # AttnBlock用のSageAttention実装
    # ... 実装内容 ...

# SageAttention検出とセットアップ
if shared.cmd_opts.use_sage_attention:
    try:
        import sageattention
        shared.sageattention_available = True
        print("[SageAttention] SageAttention detected and enabled")
    except Exception as e:
        shared.sageattention_available = False

# SageAttentionの分岐制御を追加
def attention_sage(q, k, v, heads, mask=None):
    # SageAttentionの実装
    b, _, dim_head = q.shape
    dim_head //= heads
    
    # SageAttention requires float16
    original_dtype = q.dtype
    if q.dtype != torch.float16:
        q, k, v = q.half(), k.half(), v.half()
    
    # Convert to HND layout for SageAttention: [batch_size, heads, seq_len, dim_head]
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    
    try:
        import sageattention
        global SAGE_LOGGED_THIS_GEN
        if not SAGE_LOGGED_THIS_GEN:
            print("[SageAttention] Using SageAttention 2.2 for this generation")
            SAGE_LOGGED_THIS_GEN = True
        
        # SageAttention 2.2のAPIに対応
        if hasattr(sageattention, 'sageattn'):
            out = sageattention.sageattn(q, k, v, tensor_layout='HND')
        # ... フォールバック処理 ...
        
# グローバル変数でログ制御
SAGE_LOGGED_THIS_GEN = False
FA3_LOGGED_THIS_GEN = False

# SageAttentionを最優先にチェック - UI設定も考慮
from modules.cmd_args import get_use_sage_attention
sage_should_be_used = get_use_sage_attention() or ui_sage_enabled

if sage_should_be_used:
    try:
        import sageattention
        optimized_attention = attention_sage
    except ImportError:
        sage_should_be_used = False

// SageAttention setting visibility control
function updateSageAttentionVisibility() {
    var sageAttentionSetting = gradioApp().querySelector('#setting_use_sage_attention');
    if (sageAttentionSetting) {
        var parentRow = sageAttentionSetting.closest('.form');
        if (parentRow) {
            // Check if current model is SDXL
            var isSDXL = false;
            if (typeof opts !== 'undefined' && opts.sd_model_checkpoint) {
                // This is a simple heuristic - in a real implementation, you'd want to check the actual model type
                // For now, we'll show the setting if SageAttention is available
                isSDXL = true; // Simplified for now
            }
            
            if (isSDXL) {
                parentRow.style.display = '';
            } else {
                parentRow.style.display = 'none';
            }
        }
    }
}

onAfterUiUpdate(updateSageAttentionVisibility);

from modules.forge_attention_log import reset_forge_attention_log
from ldm_patched.ldm.modules.attention import reset_sage_log, reset_fa3_log

def process_images_inner(p: StableDiffusionProcessing) -> Processed:
    try:
        from xformers.ops.fmha import reset_kernel_log
        reset_kernel_log()
    except ImportError:
        pass
    reset_forge_attention_log()  # ← 追加
    reset_sage_log()            # ← 追加
    reset_fa3_log()             # ← 追加

reset_forge_attention_log()  # ← 追加

# last item is always HTML - リストが空でないことを確認
if len(res) > 0:
    res[-1] += f"<div class='performance'><p class='time'>Time taken: <wbr><span class='measurement'>{elapsed_text}</span></p>{vram_html}{profiling_html}</div>"
else:
    # リストが空の場合は新しいHTML要素を追加
    res.append(f"<div class='performance'><p class='time'>Time taken: <wbr><span class='measurement'>{elapsed_text}</span></p>{vram_html}{profiling_html}</div>")

res[-1] += performance_html  # resが空の場合クラッシュ

python test_sageattention.py