Update src/pipeline.py

src/pipeline.py

@@ -11,14 +11,15 @@ from pipelines.models import TextToImageRequest
 from torch import Generator
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
-#from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
-#from para_attn.context_parallel import init_context_parallel_mesh
-#from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
+import torch.distributed as dist
+from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
+from para_attn.context_parallel import init_context_parallel_mesh
+from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
 from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
-#torch.cuda.set_device(dist.get_rank())
+torch.cuda.set_device(dist.get_rank())
 Pipeline = None
 ids = "slobers/Flux.1.Schnella"
 Revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
@@ -29,11 +30,11 @@ def load_pipeline() -> Pipeline:
     vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-schnell", revision="741f7c3ce8b383c54771c7003378a50191e9efe9", subfolder="vae", torch_dtype=torch.bfloat16)
     pipeline = FluxPipeline.from_pretrained(ids, revision=Revision, transformer=transformer, vae=vae, local_files_only=True, torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
- #   quantize_(pipeline.vae, int8_weight_only())
-    pipeline = apply_cache_on_pipe(pipeline, residual_diff_threshold=0.55)
- #   mesh = init_context_parallel_mesh(pipeline.device.type)
- #   parallelize_pipe(pipeline)
- #   parallelize_vae(pipeline.vae)
+    quantize_(pipeline.vae, int8_weight_only())
+ #   pipeline = apply_cache_on_pipe(pipeline, residual_diff_threshold=0.55)
+    mesh = init_context_parallel_mesh(pipeline.device.type,)
+    parallelize_pipe(pipeline, mesh=mesh,)
+    parallelize_vae(pipeline.vae, mesh=mesh._flatten())
     for _ in range(2):
         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