Update src/pipeline.py

src/pipeline.py

@@ -35,7 +35,7 @@ def load_pipeline() -> Pipeline:
     text_encoder_2 = T5EncoderModel.from_pretrained(
         "city96/t5-v1_1-xxl-encoder-bf16", revision = "1b9c856aadb864af93c1dcdc226c2774fa67bc86", torch_dtype=torch.bfloat16
     ).to(memory_format=torch.channels_last)
-
+    tinypath= os.path.join(HF_HUB_CACHE, "models--madebyollin--taef1/snapshots/5463ee684fd9131a724bea777a2f50d89b0b6b24")
     vae = AutoencoderTiny.from_pretrained(tinypath, torch_dtype=dtype)    
     path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
     model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False).to(memory_format=torch.channels_last)
@@ -43,7 +43,7 @@ def load_pipeline() -> Pipeline:
         ckpt_id,
         vae=vae,
         revision=ckpt_revision,
-        # transformer=model,
+        transformer=model,
         # text_encoder_2=text_encoder_2,
         torch_dtype=dtype,
         ).to(device)