Update src/pipeline.py

src/pipeline.py

@@ -27,47 +27,87 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
 
-# preconfigs
-import os
-os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
-os.environ["TOKENIZERS_PARALLELISM"] = "True"
-torch._dynamo.config.suppress_errors = True
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.enabled = True
 # torch.backends.cudnn.benchmark = True
+from collections import namedtuple
+import os
+Config = namedtuple('Config', ['model_id', 'revision', 'text_encoder_id', 'text_encoder_rev', 'vae_id', 'vae_rev'])
 
-# globals
 Pipeline = None
-ckpt_id = "black-forest-labs/FLUX.1-schnell"
-ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
 TinyVAE = "madebyollin/taef1"
 TinyVAE_REV = "2d552378e58c9c94201075708d7de4e1163b2689"
 
-def empty_cache():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
-    torch.cuda.reset_peak_memory_stats()
 
-def load_pipeline() -> Pipeline:
-    text_encoder_2 = T5EncoderModel.from_pretrained("manbeast3b/flux.1-schnell-full1", revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146", subfolder="text_encoder_2",torch_dtype=torch.bfloat16)
-    path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--flux.1-schnell-full1/snapshots/cb1b599b0d712b9aab2c4df3ad27b050a27ec146/transformer")
-    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
-    pipeline = FluxPipeline.from_pretrained(ckpt_id, revision=ckpt_revision, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
-    basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell-vae-kl-unst0_1_iter0/snapshots/b586f7e1125722a242c38fe963904f453095903f")
+def setup_environment():
+    os.environ.update({
+        'PYTORCH_CUDA_ALLOC_CONF': 'expandable_segments:True',
+        'TOKENIZERS_PARALLELISM': 'True',
+    })
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.enabled = True
+    torch._dynamo.config.suppress_errors = True
+
+def get_model_components(cfg):
+    text_encoder = T5EncoderModel.from_pretrained(
+        cfg.text_encoder_id,
+        revision=cfg.text_encoder_rev,
+        subfolder="text_encoder_2",
+        torch_dtype=torch.bfloat16
+    )
+    path = os.path.join(HF_HUB_CACHE, 
+                        "models--manbeast3b--flux.1-schnell-full1/snapshots/cb1b599b0d712b9aab2c4df3ad27b050a27ec146/transformer")
+    return text_encoder, path
+
+def initialize_pipeline(cfg, text_encoder, transformer_path):
+    transformer = FluxTransformer2DModel.from_pretrained(transformer_path, torch_dtype=torch.bfloat16, use_safetensors=False)
+    pipeline = FluxPipeline.from_pretrained(
+        cfg.model_id,
+        revision=cfg.revision,
+        text_encoder_2=text_encoder,
+        transformer=transformer,
+        torch_dtype=torch.bfloat16
+    )
+    basepath = os.path.join(HF_HUB_CACHE, 
+                            "models--manbeast3b--Flux.1.schnell-vae-kl-unst0_1_iter0/snapshots/b586f7e1125722a242c38fe963904f453095903f")
+
     pipeline.vae.encoder.load_state_dict(torch.load(os.path.join(basepath, "encoder.pth")), strict=False)
     pipeline.vae.decoder.load_state_dict(torch.load(os.path.join(basepath, "decoder.pth")), strict=False)
-    pipeline.to("cuda")
-    pipeline.to(memory_format=torch.channels_last)
-    for _ in range(1):
-        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.to("cuda", memory_format=torch.channels_last)
+
+def load_pipeline():
+    setup_environment()
+    torch.cuda.empty_cache()
+    cfg = Config(
+        "black-forest-labs/FLUX.1-schnell",
+        "741f7c3ce8b383c54771c7003378a50191e9efe9",
+        "manbeast3b/flux.1-schnell-full1",
+        "cb1b599b0d712b9aab2c4df3ad27b050a27ec146",
+        "manbeast3b/Flux.1.schnell-vae-kl-unst0_1_iter0",
+        "b586f7e1125722a242c38fe963904f453095903f"
+    )
+    text_encoder, transformer_path = get_model_components(cfg)
+    pipeline = initialize_pipeline(cfg, text_encoder, transformer_path)
+    warmup_ = "insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus"
+    pipeline(
+        prompt=warmup_,
+        width=1024, height=1024,
+        guidance_scale=0.0,
+        num_inference_steps=4,
+        max_sequence_length=256
+    )
+    pipeline("")
     return pipeline
 
-sample = 1
-@torch.no_grad()
-def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator) -> Image:
-    global sample
-    if not sample:
-        sample=1
-        empty_cache()
-    return pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]
+
+def infer(request, pipeline, generator):
+    with torch.no_grad():
+        result = pipeline(
+            prompt=request.prompt,
+            generator=generator,
+            guidance_scale=0.0,
+            num_inference_steps=4,
+            max_sequence_length=256,
+            height=request.height,
+            width=request.width,
+            output_type="pil"
+        )
+    return result.images[0]