Update src/pipeline.py

src/pipeline.py

@@ -17,18 +17,14 @@ os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
 
 Pipeline = None
-ids = "black-forest-labs/FLUX.1-schnell"
-Revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+ids = "slobers/Flux.1.Schnella"
 
 def load_pipeline() -> Pipeline:
-    vae = AutoencoderKL.from_pretrained(ids,revision=Revision, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16,)
-    quantize_(vae, int8_weight_only())
-    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)
-    path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
-    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False).to(memory_format=torch.channels_last)
-    pipeline = DiffusionPipeline.from_pretrained(ids, revision=Revision, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
+    transformer = FluxTransformer2DModel.from_pretrained(ids, subfolder="transformer", torch_dtype=torch.bfloat16, use_safetensors=False)
+    pipeline = FluxPipeline.from_pretrained(ids, transformer=transformer, torch_dtype=torch.bfloat16,)
     pipeline.to("cuda")
-
+    pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
+    quantize_(pipeline.vae, int8_weight_only())
     for _ in range(3):
         pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
     return pipeline
@@ -45,4 +41,4 @@ def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
         max_sequence_length=256,
         height=request.height,
         width=request.width,
-    ).images[0]
+    ).images[0]