From d7c22b3bf7b59af1c6b50f87a678eb0bee4e6335 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Thu, 22 Aug 2024 16:03:54 +0000 Subject: [PATCH] Tidy is_schnell detection logic. --- invokeai/app/invocations/flux_text_to_image.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/invokeai/app/invocations/flux_text_to_image.py b/invokeai/app/invocations/flux_text_to_image.py index 8b947b9a8d..f29b3dd309 100644 --- a/invokeai/app/invocations/flux_text_to_image.py +++ b/invokeai/app/invocations/flux_text_to_image.py @@ -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],