Update pipeline.py

pipeline.py

@@ -56,6 +56,40 @@ def prepare_timesteps(
 
 # FLUX pipeline function
 class FluxWithCFGPipeline(StableDiffusion3Pipeline):
+    def __init__(
+        self,
+        transformer: FluxTransformer2DModel,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: T5TokenizerFast,,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_3: None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            text_encoder_3=None,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            tokenizer_3=None,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 16
+        )
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = 64
+        )
+
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -208,7 +242,40 @@ class FluxWithCFGPipeline(StableDiffusion3Pipeline):
         return self.image_processor.postprocess(image, output_type=output_type)[0]
         
 class FluxWithCFGPipeline(StableDiffusion3Pipeline):
-    @torch.inference_mode()
+    def __init__(
+        self,
+        transformer: FluxTransformer2DModel,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: T5TokenizerFast,,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_3: None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            text_encoder_3=None,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            tokenizer_3=None,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 16
+        )
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = 64
+        )
+@torch.inference_mode()
     def generate_image(
         self,
         prompt: Union[str, List[str]] = None,