Implement logging for find_flux_transformer and apply_adapter methods
- Added logging configuration to control verbosity via the BOKEH_DEBUG environment variable. - Enhanced `find_flux_transformer` with detailed logging to trace the search process for FLUX transformers, including model structure and attribute checks. - Updated `apply_adapter` to log the steps taken during adapter application, including transformer verification and adapter loading. - Improved error handling with logging for better diagnostics in case of failures.
This commit is contained in:
Jesse Rabek
@@ -3,6 +3,7 @@ import numpy as np
|
||||
from typing import Optional, Tuple
|
||||
import sys
|
||||
import os
|
||||
import logging
|
||||
|
||||
# Import bokeh adapter
|
||||
from .bokeh_diffusion.adapter_flux import BokehFluxControlAdapter
|
||||
@@ -10,137 +11,339 @@ from .bokeh_diffusion.utils import color_transfer_lab
|
||||
|
||||
MIN_BOKEH, MAX_BOKEH = 0, 30
|
||||
|
||||
# Configure logging - can be controlled via BOKEH_DEBUG environment variable
|
||||
# Set BOKEH_DEBUG=1 for DEBUG level, BOKEH_DEBUG=2 for even more verbose
|
||||
_log_level = os.environ.get('BOKEH_DEBUG', '1')
|
||||
if _log_level == '2':
|
||||
logging.basicConfig(level=logging.DEBUG, format='[BokehDiffusion] %(levelname)s: %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
elif _log_level == '1':
|
||||
logging.basicConfig(level=logging.INFO, format='[BokehDiffusion] %(levelname)s: %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO)
|
||||
else:
|
||||
logging.basicConfig(level=logging.WARNING, format='[BokehDiffusion] %(levelname)s: %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.WARNING)
|
||||
|
||||
|
||||
def find_flux_transformer(model):
|
||||
"""Find the FLUX transformer in a ComfyUI model"""
|
||||
logger.info("=" * 80)
|
||||
logger.info("Starting FLUX transformer search")
|
||||
logger.info(f"Model type: {type(model)}")
|
||||
logger.info(f"Model ID: {id(model)}")
|
||||
|
||||
# Helper function to check if an object is a FLUX transformer
|
||||
def is_flux_transformer(obj):
|
||||
def is_flux_transformer(obj, obj_name="unknown"):
|
||||
if obj is None:
|
||||
logger.debug(f" [is_flux_transformer] {obj_name}: Object is None")
|
||||
return False
|
||||
obj_type = str(type(obj))
|
||||
logger.debug(f" [is_flux_transformer] {obj_name}: Type = {obj_type}")
|
||||
|
||||
# Check for FluxTransformer2DModel class name
|
||||
if 'FluxTransformer2DModel' in obj_type:
|
||||
logger.info(f" [is_flux_transformer] {obj_name}: ✓ MATCH - FluxTransformer2DModel detected!")
|
||||
return True
|
||||
|
||||
# Check for attn_processors attribute (key indicator of FLUX transformer)
|
||||
if hasattr(obj, 'attn_processors'):
|
||||
has_attn_procs = hasattr(obj, 'attn_processors')
|
||||
has_set_attn = hasattr(obj, 'set_attn_processor')
|
||||
logger.debug(f" [is_flux_transformer] {obj_name}: has attn_processors={has_attn_procs}, has set_attn_processor={has_set_attn}")
|
||||
|
||||
if has_attn_procs:
|
||||
# Additional check: FLUX transformers have set_attn_processor method
|
||||
if hasattr(obj, 'set_attn_processor'):
|
||||
if has_set_attn:
|
||||
logger.info(f" [is_flux_transformer] {obj_name}: ✓ MATCH - Has attn_processors and set_attn_processor!")
|
||||
return True
|
||||
else:
|
||||
logger.debug(f" [is_flux_transformer] {obj_name}: Has attn_processors but missing set_attn_processor")
|
||||
|
||||
logger.debug(f" [is_flux_transformer] {obj_name}: ✗ Not a transformer")
|
||||
return False
|
||||
|
||||
# Try different possible locations for the transformer
|
||||
# Most common: model.model.diffusion_model (ComfyUI ModelPatcher structure)
|
||||
logger.info("Checking model.model structure...")
|
||||
if hasattr(model, 'model'):
|
||||
logger.info(f" model.model exists: {type(model.model)}")
|
||||
logger.debug(f" model.model attributes: {[a for a in dir(model.model) if not a.startswith('__')][:20]}")
|
||||
|
||||
if hasattr(model.model, 'diffusion_model'):
|
||||
logger.info(" Found model.model.diffusion_model")
|
||||
flux_model = model.model.diffusion_model
|
||||
flux_type = str(type(flux_model))
|
||||
logger.info(f" flux_model type: {flux_type}")
|
||||
|
||||
# Check if it's a ComfyUI Flux wrapper (comfy.ldm.flux.model.Flux)
|
||||
# In this case, the actual transformer is usually in flux_model.transformer
|
||||
if 'comfy.ldm.flux' in flux_type or 'comfy.model_base.Flux' in flux_type:
|
||||
logger.info(" ✓ Detected ComfyUI Flux wrapper - searching for transformer inside...")
|
||||
|
||||
# The actual FluxTransformer2DModel is typically in the 'transformer' attribute
|
||||
logger.info(" Checking flux_model.transformer...")
|
||||
if hasattr(flux_model, 'transformer'):
|
||||
logger.info(" flux_model.transformer exists!")
|
||||
transformer = flux_model.transformer
|
||||
if is_flux_transformer(transformer):
|
||||
logger.info(f" transformer type: {type(transformer)}")
|
||||
if is_flux_transformer(transformer, "flux_model.transformer"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at flux_model.transformer!")
|
||||
return transformer
|
||||
# Also check other common locations
|
||||
for attr_name in ['model', '_model']:
|
||||
else:
|
||||
logger.warning(" flux_model.transformer exists but is_flux_transformer returned False")
|
||||
else:
|
||||
logger.warning(" flux_model.transformer attribute NOT FOUND")
|
||||
|
||||
# Check all direct attributes (including private ones)
|
||||
logger.info(" Checking direct attributes: model, _model, transformer, _transformer...")
|
||||
for attr_name in ['model', '_model', 'transformer', '_transformer']:
|
||||
logger.debug(f" Checking attribute: {attr_name}")
|
||||
if hasattr(flux_model, attr_name):
|
||||
inner_obj = getattr(flux_model, attr_name)
|
||||
if inner_obj is not None and is_flux_transformer(inner_obj):
|
||||
return inner_obj
|
||||
logger.info(f" {attr_name} exists!")
|
||||
try:
|
||||
inner_obj = getattr(flux_model, attr_name)
|
||||
logger.debug(f" {attr_name} type: {type(inner_obj)}")
|
||||
if inner_obj is not None:
|
||||
if is_flux_transformer(inner_obj, f"flux_model.{attr_name}"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER at flux_model.{attr_name}!")
|
||||
return inner_obj
|
||||
# Recursively check if it contains the transformer
|
||||
if hasattr(inner_obj, 'attn_processors') or 'FluxTransformer' in str(type(inner_obj)):
|
||||
logger.debug(f" {attr_name} has attn_processors or FluxTransformer in type")
|
||||
# Check nested
|
||||
if hasattr(inner_obj, 'transformer'):
|
||||
logger.info(f" {attr_name}.transformer exists!")
|
||||
nested = getattr(inner_obj, 'transformer')
|
||||
if is_flux_transformer(nested, f"flux_model.{attr_name}.transformer"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER at flux_model.{attr_name}.transformer!")
|
||||
return nested
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logger.debug(f" Error accessing {attr_name}: {e}")
|
||||
continue
|
||||
else:
|
||||
logger.debug(f" {attr_name} does not exist")
|
||||
|
||||
# Search through all named modules/children to find the transformer
|
||||
# ComfyUI might store it as a submodule
|
||||
logger.info(" Searching through named_modules...")
|
||||
transformer_candidates = []
|
||||
attn_processor_candidates = []
|
||||
try:
|
||||
module_count = 0
|
||||
for name, module in flux_model.named_modules():
|
||||
if is_flux_transformer(module):
|
||||
return module
|
||||
except (AttributeError, RuntimeError, TypeError):
|
||||
module_count += 1
|
||||
module_type = str(type(module))
|
||||
# Check if it's a FluxTransformer2DModel
|
||||
if 'FluxTransformer2DModel' in module_type:
|
||||
logger.debug(f" Found FluxTransformer2DModel candidate: {name}")
|
||||
transformer_candidates.append((name, module_type))
|
||||
if is_flux_transformer(module, f"named_modules[{name}]"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER in named_modules: {name} ({type(module)})!")
|
||||
return module
|
||||
# Check if it has attn_processors
|
||||
elif hasattr(module, 'attn_processors') and hasattr(module, 'set_attn_processor'):
|
||||
logger.debug(f" Found attn_processor candidate: {name}")
|
||||
attn_processor_candidates.append((name, module_type))
|
||||
if is_flux_transformer(module, f"named_modules[{name}]"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER-LIKE in named_modules: {name} ({type(module)})!")
|
||||
return module
|
||||
|
||||
logger.info(f" Searched {module_count} modules total")
|
||||
logger.info(f" Found {len(transformer_candidates)} FluxTransformer2DModel candidates")
|
||||
logger.info(f" Found {len(attn_processor_candidates)} attn_processor candidates")
|
||||
|
||||
# Print candidate modules for debugging
|
||||
if transformer_candidates:
|
||||
logger.info(" FluxTransformer2DModel candidates:")
|
||||
for name, mod_type in transformer_candidates[:20]: # Show first 20
|
||||
logger.info(f" {name}: {mod_type}")
|
||||
|
||||
if attn_processor_candidates:
|
||||
logger.info(" attn_processor candidates:")
|
||||
for name, mod_type in attn_processor_candidates[:20]: # Show first 20
|
||||
logger.info(f" {name}: {mod_type}")
|
||||
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logger.error(f" Error searching named_modules: {e}", exc_info=True)
|
||||
pass
|
||||
|
||||
# Also check direct children
|
||||
logger.info(" Checking direct children...")
|
||||
try:
|
||||
child_count = 0
|
||||
for child in flux_model.children():
|
||||
if is_flux_transformer(child):
|
||||
child_count += 1
|
||||
logger.debug(f" Child {child_count}: {type(child)}")
|
||||
if is_flux_transformer(child, f"children[{child_count}]"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER in children: {type(child)}!")
|
||||
return child
|
||||
except (AttributeError, RuntimeError, TypeError):
|
||||
logger.info(f" Checked {child_count} direct children")
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logger.error(f" Error searching children: {e}", exc_info=True)
|
||||
pass
|
||||
|
||||
# Try searching through all attributes more thoroughly
|
||||
# Check if any attribute contains FluxTransformer2DModel
|
||||
logger.info(" Searching through all private attributes...")
|
||||
try:
|
||||
all_attrs = [a for a in dir(flux_model) if a.startswith('_') and not a.startswith('__')]
|
||||
logger.debug(f" Found {len(all_attrs)} private attributes to check")
|
||||
checked_count = 0
|
||||
for attr_name in all_attrs:
|
||||
checked_count += 1
|
||||
if checked_count % 10 == 0:
|
||||
logger.debug(f" Checked {checked_count}/{len(all_attrs)} attributes...")
|
||||
try:
|
||||
attr_obj = getattr(flux_model, attr_name, None)
|
||||
if attr_obj is not None:
|
||||
attr_type = str(type(attr_obj))
|
||||
# Check if it's a FluxTransformer2DModel
|
||||
if 'FluxTransformer2DModel' in attr_type:
|
||||
logger.info(f" Found FluxTransformer2DModel in attribute: {attr_name} ({attr_type})")
|
||||
if is_flux_transformer(attr_obj, f"attr[{attr_name}]"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER in attribute: {attr_name}!")
|
||||
return attr_obj
|
||||
# Check if it has attn_processors
|
||||
if hasattr(attr_obj, 'attn_processors') and hasattr(attr_obj, 'set_attn_processor'):
|
||||
logger.info(f" Found transformer-like object in attribute: {attr_name} ({attr_type})")
|
||||
if is_flux_transformer(attr_obj, f"attr[{attr_name}]"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER in attribute: {attr_name}!")
|
||||
return attr_obj
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logger.debug(f" Error accessing {attr_name}: {e}")
|
||||
continue
|
||||
logger.info(f" Checked {checked_count} private attributes")
|
||||
except Exception as e:
|
||||
logger.error(f" Error searching attributes: {e}", exc_info=True)
|
||||
pass
|
||||
else:
|
||||
# Not a ComfyUI wrapper, check if it's the transformer directly
|
||||
if is_flux_transformer(flux_model):
|
||||
logger.info(" Not a ComfyUI wrapper, checking if flux_model is transformer directly...")
|
||||
if is_flux_transformer(flux_model, "flux_model (direct)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER - flux_model is transformer directly!")
|
||||
return flux_model
|
||||
else:
|
||||
logger.warning(" model.model.diffusion_model does not exist")
|
||||
|
||||
logger.info(" Checking model.model.transformer...")
|
||||
if hasattr(model.model, 'transformer'):
|
||||
logger.info(" model.model.transformer exists!")
|
||||
transformer = model.model.transformer
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "model.model.transformer"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at model.model.transformer!")
|
||||
return transformer
|
||||
else:
|
||||
logger.debug(" model.model.transformer does not exist")
|
||||
|
||||
# Sometimes the model.model itself might be the transformer
|
||||
if is_flux_transformer(model.model):
|
||||
logger.info(" Checking if model.model is transformer directly...")
|
||||
if is_flux_transformer(model.model, "model.model (direct)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER - model.model is transformer directly!")
|
||||
return model.model
|
||||
else:
|
||||
logger.warning(" model.model does not exist")
|
||||
|
||||
# Check for private _model attribute (some wrappers use this)
|
||||
logger.info("Checking model._model...")
|
||||
if hasattr(model, '_model'):
|
||||
logger.info(" model._model exists!")
|
||||
try:
|
||||
inner_model = model._model
|
||||
if inner_model is not None:
|
||||
logger.info(f" _model type: {type(inner_model)}")
|
||||
if hasattr(inner_model, 'diffusion_model'):
|
||||
logger.info(" _model.diffusion_model exists!")
|
||||
transformer = inner_model.diffusion_model
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "model._model.diffusion_model"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at model._model.diffusion_model!")
|
||||
return transformer
|
||||
if hasattr(inner_model, 'transformer'):
|
||||
logger.info(" _model.transformer exists!")
|
||||
transformer = inner_model.transformer
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "model._model.transformer"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at model._model.transformer!")
|
||||
return transformer
|
||||
if is_flux_transformer(inner_model):
|
||||
if is_flux_transformer(inner_model, "model._model (direct)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER - model._model is transformer directly!")
|
||||
return inner_model
|
||||
except (AttributeError, RuntimeError, TypeError):
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logger.debug(f" Error accessing _model: {e}")
|
||||
pass
|
||||
else:
|
||||
logger.debug(" model._model does not exist")
|
||||
|
||||
# Check if model has a get_model() method (some ComfyUI wrappers)
|
||||
logger.info("Checking model.get_model()...")
|
||||
if hasattr(model, 'get_model'):
|
||||
logger.info(" model.get_model() exists!")
|
||||
try:
|
||||
inner_model = model.get_model()
|
||||
logger.info(f" get_model() returned: {type(inner_model)}")
|
||||
if inner_model is not None:
|
||||
if hasattr(inner_model, 'diffusion_model'):
|
||||
logger.info(" inner_model.diffusion_model exists!")
|
||||
transformer = inner_model.diffusion_model
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "get_model().diffusion_model"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at get_model().diffusion_model!")
|
||||
return transformer
|
||||
if hasattr(inner_model, 'transformer'):
|
||||
logger.info(" inner_model.transformer exists!")
|
||||
transformer = inner_model.transformer
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "get_model().transformer"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at get_model().transformer!")
|
||||
return transformer
|
||||
if is_flux_transformer(inner_model):
|
||||
if is_flux_transformer(inner_model, "get_model() (direct)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER - get_model() result is transformer directly!")
|
||||
return inner_model
|
||||
except (AttributeError, RuntimeError, TypeError):
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logger.debug(f" Error calling get_model(): {e}")
|
||||
pass
|
||||
else:
|
||||
logger.debug(" model.get_model() does not exist")
|
||||
|
||||
# Direct access: model.diffusion_model
|
||||
logger.info("Checking model.diffusion_model (direct)...")
|
||||
if hasattr(model, 'diffusion_model'):
|
||||
logger.info(" model.diffusion_model exists!")
|
||||
transformer = model.diffusion_model
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "model.diffusion_model (direct)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at model.diffusion_model!")
|
||||
return transformer
|
||||
else:
|
||||
logger.debug(" model.diffusion_model does not exist")
|
||||
|
||||
# Direct access: model.transformer
|
||||
logger.info("Checking model.transformer (direct)...")
|
||||
if hasattr(model, 'transformer'):
|
||||
logger.info(" model.transformer exists!")
|
||||
transformer = model.transformer
|
||||
if is_flux_transformer(transformer):
|
||||
if is_flux_transformer(transformer, "model.transformer (direct)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER at model.transformer!")
|
||||
return transformer
|
||||
else:
|
||||
logger.debug(" model.transformer does not exist")
|
||||
|
||||
# Last resort: check if model itself is the transformer
|
||||
if is_flux_transformer(model):
|
||||
logger.info("Checking if model itself is transformer...")
|
||||
if is_flux_transformer(model, "model (self)"):
|
||||
logger.info(" ✓ FOUND TRANSFORMER - model itself is transformer!")
|
||||
return model
|
||||
|
||||
# Recursive search as final fallback
|
||||
def search_recursive(obj, depth=0, max_depth=3, visited=None):
|
||||
logger.info("Starting recursive search as final fallback...")
|
||||
def search_recursive(obj, depth=0, max_depth=3, visited=None, path=""):
|
||||
if visited is None:
|
||||
visited = set()
|
||||
if depth > max_depth or id(obj) in visited:
|
||||
return None
|
||||
visited.add(id(obj))
|
||||
logger.debug(f" [recursive] depth={depth}, path={path}, type={type(obj)}")
|
||||
|
||||
# Check if current object is the transformer
|
||||
if is_flux_transformer(obj):
|
||||
if is_flux_transformer(obj, f"recursive[{path}]"):
|
||||
logger.info(f" ✓ FOUND TRANSFORMER in recursive search at: {path}!")
|
||||
return obj
|
||||
|
||||
# Search in common attribute names
|
||||
@@ -149,7 +352,8 @@ def find_flux_transformer(model):
|
||||
try:
|
||||
attr_obj = getattr(obj, attr_name)
|
||||
if attr_obj is not None and id(attr_obj) not in visited:
|
||||
result = search_recursive(attr_obj, depth + 1, max_depth, visited)
|
||||
new_path = f"{path}.{attr_name}" if path else attr_name
|
||||
result = search_recursive(attr_obj, depth + 1, max_depth, visited, new_path)
|
||||
if result is not None:
|
||||
return result
|
||||
except (AttributeError, RuntimeError, TypeError):
|
||||
@@ -160,7 +364,8 @@ def find_flux_transformer(model):
|
||||
try:
|
||||
inner_model = obj.get_model()
|
||||
if inner_model is not None and id(inner_model) not in visited:
|
||||
result = search_recursive(inner_model, depth + 1, max_depth, visited)
|
||||
new_path = f"{path}.get_model()" if path else "get_model()"
|
||||
result = search_recursive(inner_model, depth + 1, max_depth, visited, new_path)
|
||||
if result is not None:
|
||||
return result
|
||||
except (AttributeError, RuntimeError, TypeError):
|
||||
@@ -168,37 +373,50 @@ def find_flux_transformer(model):
|
||||
|
||||
return None
|
||||
|
||||
transformer = search_recursive(model)
|
||||
transformer = search_recursive(model, path="model")
|
||||
if transformer is not None:
|
||||
logger.info(" ✓ FOUND TRANSFORMER via recursive search!")
|
||||
return transformer
|
||||
|
||||
# Debug: print model structure to help diagnose
|
||||
print("DEBUG: Model structure:")
|
||||
print(f" Model type: {type(model)}")
|
||||
# Final debug summary
|
||||
logger.error("=" * 80)
|
||||
logger.error("FAILED TO FIND FLUX TRANSFORMER")
|
||||
logger.error("=" * 80)
|
||||
logger.error("Model structure summary:")
|
||||
logger.error(f" Model type: {type(model)}")
|
||||
logger.error(f" Model attributes: {[a for a in dir(model) if not a.startswith('__')][:30]}")
|
||||
if hasattr(model, 'model'):
|
||||
print(f" model.model type: {type(model.model)}")
|
||||
logger.error(f" model.model type: {type(model.model)}")
|
||||
logger.error(f" model.model attributes: {[a for a in dir(model.model) if not a.startswith('__')][:30]}")
|
||||
if hasattr(model.model, 'diffusion_model'):
|
||||
flux_model = model.model.diffusion_model
|
||||
flux_type = str(type(flux_model))
|
||||
print(f" model.model.diffusion_model type: {flux_type}")
|
||||
print(f" Has attn_processors: {hasattr(flux_model, 'attn_processors')}")
|
||||
logger.error(f" model.model.diffusion_model type: {flux_type}")
|
||||
logger.error(f" Has attn_processors: {hasattr(flux_model, 'attn_processors')}")
|
||||
logger.error(f" Has set_attn_processor: {hasattr(flux_model, 'set_attn_processor')}")
|
||||
|
||||
# Check if it's a ComfyUI Flux wrapper
|
||||
if 'comfy.ldm.flux' in flux_type or 'comfy.model_base.Flux' in flux_type:
|
||||
print(" Detected ComfyUI Flux wrapper - searching for transformer...")
|
||||
logger.error(" Detected ComfyUI Flux wrapper")
|
||||
# Check for transformer attribute
|
||||
if hasattr(flux_model, 'transformer'):
|
||||
transformer = flux_model.transformer
|
||||
print(f" Found flux_model.transformer: {type(transformer)}")
|
||||
print(f" Has attn_processors: {hasattr(transformer, 'attn_processors')}")
|
||||
print(f" Has set_attn_processor: {hasattr(transformer, 'set_attn_processor')}")
|
||||
logger.error(f" flux_model.transformer exists: {type(transformer)}")
|
||||
logger.error(f" Has attn_processors: {hasattr(transformer, 'attn_processors')}")
|
||||
logger.error(f" Has set_attn_processor: {hasattr(transformer, 'set_attn_processor')}")
|
||||
else:
|
||||
print(" flux_model.transformer attribute NOT FOUND")
|
||||
logger.error(" flux_model.transformer attribute NOT FOUND")
|
||||
# List all attributes that might be relevant
|
||||
relevant_attrs = [attr for attr in dir(flux_model)
|
||||
if not attr.startswith('__') and
|
||||
not isinstance(getattr(flux_model, attr, None), (str, int, float, bool, list, dict, tuple))]
|
||||
print(f" Relevant attributes: {relevant_attrs[:20]}") # Show first 20
|
||||
all_attrs = [attr for attr in dir(flux_model) if not attr.startswith('__')]
|
||||
logger.error(f" All non-private attributes ({len(all_attrs)} total):")
|
||||
for attr in all_attrs[:50]: # Show first 50
|
||||
try:
|
||||
attr_obj = getattr(flux_model, attr, None)
|
||||
if attr_obj is not None and not isinstance(attr_obj, (str, int, float, bool, list, dict, tuple)):
|
||||
logger.error(f" {attr}: {type(attr_obj)}")
|
||||
except:
|
||||
pass
|
||||
logger.error("=" * 80)
|
||||
|
||||
return None
|
||||
|
||||
@@ -266,47 +484,103 @@ class ApplyBokehAdapter:
|
||||
|
||||
def apply_adapter(self, model, adapter):
|
||||
"""Apply the bokeh adapter to the model's transformer"""
|
||||
logger.info("=" * 80)
|
||||
logger.info("ApplyBokehAdapter.apply_adapter() called")
|
||||
logger.info(f"Adapter config: {adapter}")
|
||||
try:
|
||||
# Find the transformer
|
||||
logger.info("Step 1: Finding FLUX transformer...")
|
||||
transformer = find_flux_transformer(model)
|
||||
|
||||
if transformer is None:
|
||||
logger.error("Transformer is None - search failed!")
|
||||
raise Exception("Could not find FLUX transformer in model. Make sure you're using a FLUX model.")
|
||||
|
||||
logger.info(f"✓ Transformer found: {type(transformer)}")
|
||||
logger.info(f" Transformer ID: {id(transformer)}")
|
||||
|
||||
# Check if transformer has attn_processors (required for adapter)
|
||||
logger.info("Step 2: Verifying transformer has attn_processors...")
|
||||
if not hasattr(transformer, 'attn_processors'):
|
||||
logger.error("Transformer does not have attn_processors attribute!")
|
||||
raise Exception("Model transformer does not have attn_processors. This may not be a FLUX model.")
|
||||
|
||||
logger.info("✓ Transformer has attn_processors")
|
||||
logger.debug(f" attn_processors type: {type(transformer.attn_processors)}")
|
||||
if hasattr(transformer.attn_processors, '__len__'):
|
||||
logger.debug(f" Number of processors: {len(transformer.attn_processors)}")
|
||||
|
||||
# Check for set_attn_processor method
|
||||
if not hasattr(transformer, 'set_attn_processor'):
|
||||
logger.error("Transformer does not have set_attn_processor method!")
|
||||
raise Exception("Model transformer does not have set_attn_processor method. This may not be a FLUX model.")
|
||||
logger.info("✓ Transformer has set_attn_processor method")
|
||||
|
||||
# Load the adapter if not already loaded
|
||||
adapter_name = adapter["adapter_name"]
|
||||
logger.info(f"Step 3: Loading adapter from {adapter_name}...")
|
||||
if adapter["adapter"] is None:
|
||||
print(f"Loading Bokeh adapter from {adapter_name}...")
|
||||
adapter_obj = BokehFluxControlAdapter.from_pretrained(
|
||||
adapter_name,
|
||||
base_model=transformer,
|
||||
)
|
||||
# Move to same device as transformer
|
||||
device = next(transformer.parameters()).device
|
||||
dtype = next(transformer.parameters()).dtype
|
||||
adapter_obj = adapter_obj.to(device=device, dtype=dtype)
|
||||
adapter["adapter"] = adapter_obj
|
||||
print("Bokeh adapter loaded successfully.")
|
||||
logger.info(f" Adapter not yet loaded, loading now...")
|
||||
try:
|
||||
adapter_obj = BokehFluxControlAdapter.from_pretrained(
|
||||
adapter_name,
|
||||
base_model=transformer,
|
||||
)
|
||||
logger.info(" ✓ Adapter loaded from HuggingFace")
|
||||
|
||||
# Move to same device as transformer
|
||||
logger.info("Step 4: Moving adapter to same device/dtype as transformer...")
|
||||
try:
|
||||
device = next(transformer.parameters()).device
|
||||
dtype = next(transformer.parameters()).dtype
|
||||
logger.info(f" Transformer device: {device}, dtype: {dtype}")
|
||||
adapter_obj = adapter_obj.to(device=device, dtype=dtype)
|
||||
logger.info(" ✓ Adapter moved to transformer device/dtype")
|
||||
except Exception as e:
|
||||
logger.warning(f" Error moving adapter to device: {e}")
|
||||
logger.info(" Continuing anyway...")
|
||||
|
||||
adapter["adapter"] = adapter_obj
|
||||
logger.info("✓ Bokeh adapter loaded and configured successfully")
|
||||
except Exception as e:
|
||||
logger.error(f" Failed to load adapter: {e}", exc_info=True)
|
||||
raise
|
||||
else:
|
||||
logger.info(" Adapter already loaded, reusing existing instance")
|
||||
|
||||
adapter_obj = adapter["adapter"]
|
||||
logger.info(f" Adapter object: {type(adapter_obj)}")
|
||||
|
||||
# Set scales
|
||||
if adapter["bokeh_scale"] != 1.0:
|
||||
adapter_obj.set_bokeh_scale(adapter["bokeh_scale"])
|
||||
if adapter["lora_scale"] != 1.0:
|
||||
adapter_obj.set_lora_scale(adapter["lora_scale"])
|
||||
logger.info("Step 5: Setting adapter scales...")
|
||||
bokeh_scale = adapter.get("bokeh_scale", 1.0)
|
||||
lora_scale = adapter.get("lora_scale", 1.0)
|
||||
logger.info(f" bokeh_scale: {bokeh_scale}, lora_scale: {lora_scale}")
|
||||
|
||||
if bokeh_scale != 1.0:
|
||||
logger.info(f" Setting bokeh_scale to {bokeh_scale}")
|
||||
adapter_obj.set_bokeh_scale(bokeh_scale)
|
||||
if lora_scale != 1.0:
|
||||
logger.info(f" Setting lora_scale to {lora_scale}")
|
||||
adapter_obj.set_lora_scale(lora_scale)
|
||||
|
||||
# Store adapter in model for use during sampling
|
||||
logger.info("Step 6: Storing adapter in model...")
|
||||
model.bokeh_adapter = adapter_obj
|
||||
model.bokeh_adapter_config = adapter
|
||||
model._bokeh_transformer = transformer
|
||||
logger.info("✓ Adapter stored in model")
|
||||
logger.info("=" * 80)
|
||||
logger.info("✓ Bokeh adapter applied successfully!")
|
||||
logger.info("=" * 80)
|
||||
|
||||
return (model,)
|
||||
except Exception as e:
|
||||
logger.error("=" * 80)
|
||||
logger.error("FAILED TO APPLY BOKEH ADAPTER")
|
||||
logger.error("=" * 80)
|
||||
logger.error(f"Error: {e}", exc_info=True)
|
||||
logger.error("=" * 80)
|
||||
raise Exception(f"Failed to apply bokeh adapter: {e}")
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user