Tidy is_schnell detection logic.

This commit is contained in:
Ryan Dick 2024-08-22 16:03:54 +00:00 committed by Brandon
parent 185f2a395f
commit d7c22b3bf7

View File

@ -17,7 +17,6 @@ from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.flux.model import Flux
from invokeai.backend.flux.modules.autoencoder import AutoEncoder
from invokeai.backend.flux.sampling import denoise, get_noise, get_schedule, unpack
from invokeai.backend.model_manager.config import CheckpointConfigBase
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
from invokeai.backend.util.devices import TorchDevice
@ -90,12 +89,8 @@ class FluxTextToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
img, img_ids = self._prepare_latent_img_patches(x)
# HACK(ryand): Find a better way to determine if this is a schnell model or not.
is_schnell = (
"schnell" in transformer_info.config.config_path
if transformer_info.config and isinstance(transformer_info.config, CheckpointConfigBase)
else ""
)
is_schnell = "schnell" in transformer_info.config.config_path
timesteps = get_schedule(
num_steps=self.num_steps,
image_seq_len=img.shape[1],