From fb5db32bb0481c611268664d232e2fd1536e9590 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Mon, 19 Aug 2024 22:12:12 +0000 Subject: [PATCH] Make it run (with artifacts). --- invokeai/app/invocations/flux_text_to_image.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/invokeai/app/invocations/flux_text_to_image.py b/invokeai/app/invocations/flux_text_to_image.py index 19829c47a4..229c1fdf46 100644 --- a/invokeai/app/invocations/flux_text_to_image.py +++ b/invokeai/app/invocations/flux_text_to_image.py @@ -118,7 +118,7 @@ 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 = "shnell" in str(flux_transformer_path) + is_schnell = "schnell" in str(flux_transformer_path) timesteps = get_schedule( num_steps=self.num_steps, image_seq_len=img.shape[1], @@ -174,6 +174,7 @@ class FluxTextToImageInvocation(BaseInvocation, WithMetadata, WithBoard): img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None] img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :] img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) + img_ids = img_ids.to(latent_img.device) return img, img_ids @@ -239,13 +240,13 @@ class FluxTextToImageInvocation(BaseInvocation, WithMetadata, WithBoard): else: raise ValueError(f"Unsupported quantization type: {self.quantization_type}") - assert isinstance(model, FluxTransformer2DModel) + assert isinstance(model, Flux) return model @staticmethod def _load_flux_vae(path: Path) -> AutoEncoder: # TODO(ryand): Determine if this is a schnell model or a dev model and load the appropriate config. - ae_params = flux_configs["flux1-schnell"].ae_params + ae_params = flux_configs["flux-schnell"].ae_params with accelerate.init_empty_weights(): ae = AutoEncoder(ae_params)