Update src/pipeline.py

src/pipeline.py

@@ -159,177 +159,138 @@
 # are_two_tensors_similar_max_diff = TensorComparator.max_diff_comparison
 # empty_cache = MemoryManager.empty_cache
 
-
-from __future__ import annotations
 import os
+import gc
 import torch
-import functools
-from enum import Enum, auto
-from contextlib import contextmanager
-from typing import Protocol, TypeVar, Generic, Callable, Any
-from dataclasses import dataclass, field
 from PIL.Image import Image
-from torch import Generator
+from dataclasses import dataclass
 from diffusers import DiffusionPipeline, AutoencoderTiny, FluxTransformer2DModel
 from transformers import T5EncoderModel
 from huggingface_hub.constants import HF_HUB_CACHE
 from torchao.quantization import quantize_, int8_weight_only, float8_weight_only
 from first_block_cache.diffusers_adapters import apply_cache_on_pipe
 from pipelines.models import TextToImageRequest
+from torch import Generator
 
-T = TypeVar('T')
-
-class ModelComponent(Protocol):
-    def to(self, *args, **kwargs) -> Any: ...
-
-class ModelState(Enum):
-    INITIALIZED = auto()
-    LOADED = auto()
-    OPTIMIZED = auto()
-    READY = auto()
-
-class ResourceMonitor:
-    """Monitors and manages system resources."""
-    
-    @contextmanager
-    def monitor_memory(self, threshold_mb: int = 1000):
-        initial_memory = torch.cuda.memory_allocated() / 1024**2
-        yield
-        final_memory = torch.cuda.memory_allocated() / 1024**2
-        if final_memory - initial_memory > threshold_mb:
-            torch.cuda.empty_cache()
-
+# Configuration settings using a dataclass for clarity
 @dataclass
-class ModelRegistry(Generic[T]):
-    """Type-safe registry for model components."""
-    _components: dict[str, T] = field(default_factory=dict)
-    
-    def register(self, name: str, component: T) -> None:
-        self._components[name] = component
-    
-    def get(self, name: str) -> T:
-        return self._components[name]
-    
-    def __iter__(self):
-        return iter(self._components.values())
-
-class PipelineBuilder:
-    """Fluent builder for pipeline construction."""
-    
-    def __init__(self):
-        self.config = {
-            "model_id": "black-forest-labs/FLUX.1-schnell",
-            "revision": "741f7c3ce8b383c54771c7003378a50191e9efe9",
-            "device": "cuda",
-            "dtype": torch.bfloat16
-        }
-        self.registry = ModelRegistry[ModelComponent]()
-        self.state = ModelState.INITIALIZED
-        self.monitor = ResourceMonitor()
-
-    def with_torch_settings(self) -> PipelineBuilder:
-        torch.backends.cuda.matmul.allow_tf32 = True
-        torch.backends.cudnn.enabled = True
-        torch.backends.cudnn.benchmark = True
-        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = "expandable_segments:True"
-        return self
-
-    def load_text_encoder(self) -> PipelineBuilder:
-        with self.monitor.monitor_memory():
-            encoder = T5EncoderModel.from_pretrained(
-                "city96/t5-v1_1-xxl-encoder-bf16",
-                revision="1b9c856aadb864af93c1dcdc226c2774fa67bc86",
-                torch_dtype=self.config["dtype"]
-            ).to(memory_format=torch.channels_last)
-            self.registry.register("text_encoder", encoder)
-        return self
-
-    def load_vae(self) -> PipelineBuilder:
-        with self.monitor.monitor_memory():
-            vae = AutoencoderTiny.from_pretrained(
-                "RobertML/FLUX.1-schnell-vae_e3m2",
-                revision="da0d2cd7815792fb40d084dbd8ed32b63f153d8d",
-                torch_dtype=self.config["dtype"]
-            )
-            self.registry.register("vae", vae)
-        return self
-
-    def load_transformer(self) -> PipelineBuilder:
-        with self.monitor.monitor_memory():
-            path = os.path.join(
-                HF_HUB_CACHE,
-                "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a"
-            )
-            transformer = FluxTransformer2DModel.from_pretrained(
-                path,
-                torch_dtype=self.config["dtype"],
-                use_safetensors=False
-            ).to(memory_format=torch.channels_last)
-            self.registry.register("transformer", transformer)
-        return self
-
-    def optimize(self, pipeline: DiffusionPipeline) -> PipelineBuilder:
-        with self.monitor.monitor_memory():
-            pipeline.to(memory_format=torch.channels_last)
-            pipeline.vae = torch.compile(pipeline.vae, mode="max-autotune")
-            quantize_(pipeline.vae, int8_weight_only())
-            quantize_(pipeline.vae, float8_weight_only())
-            apply_cache_on_pipe(pipeline)
-        return self
-
-    def warmup(self, pipeline: DiffusionPipeline) -> PipelineBuilder:
-        with torch.no_grad(), self.monitor.monitor_memory():
-            for _ in range(3):
-                pipeline(prompt=" ")
-        return self
-
-    def build(self) -> DiffusionPipeline:
-        pipeline = DiffusionPipeline.from_pretrained(
-            self.config["model_id"],
-            vae=self.registry.get("vae"),
-            revision=self.config["revision"],
-            transformer=self.registry.get("transformer"),
-            text_encoder_2=self.registry.get("text_encoder"),
-            torch_dtype=self.config["dtype"],
-        ).to(self.config["device"])
-        
-        self.optimize(pipeline)
-        self.warmup(pipeline)
-        self.state = ModelState.READY
-        return pipeline
-
-class InferenceContext:
-    """Context manager for inference operations."""
-    
-    def __init__(self, pipeline: DiffusionPipeline):
-        self.pipeline = pipeline
-        self.monitor = ResourceMonitor()
-    
-    @contextmanager
-    def inference_mode(self):
-        with torch.no_grad(), self.monitor.monitor_memory():
-            yield self.pipeline
-
-def load() -> DiffusionPipeline:
-    """Build and configure the pipeline using the fluent builder pattern."""
-    return (PipelineBuilder()
-            .with_torch_settings()
-            .load_text_encoder()
-            .load_vae()
-            .load_transformer()
-            .build())
-
+class Config:
+    CKPT_ID: str = "black-forest-labs/FLUX.1-schnell"
+    CKPT_REVISION: str = "741f7c3ce8b383c54771c7003378a50191e9efe9"
+    DEVICE: str = "cuda"
+    DTYPE = torch.bfloat16
+    PYTORCH_CUDA_ALLOC_CONF: str = "expandable_segments:True"
+
+def _initialize_environment():
+    """Set up PyTorch and CUDA environment variables for optimal performance."""
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.benchmark = True
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = Config.PYTORCH_CUDA_ALLOC_CONF
+
+def _clear_gpu_memory():
+    """Free up GPU memory to prevent memory-related issues."""
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+
+def _load_text_encoder_model():
+    """Load the text encoder model with specified configuration."""
+    return T5EncoderModel.from_pretrained(
+        "city96/t5-v1_1-xxl-encoder-bf16",
+        revision="1b9c856aadb864af93c1dcdc226c2774fa67bc86",
+        torch_dtype=Config.DTYPE
+    ).to(memory_format=torch.channels_last)
+
+def _load_vae_model():
+    """Load the variational autoencoder (VAE) model with specified configuration."""
+    return AutoencoderTiny.from_pretrained(
+        "RobertML/FLUX.1-schnell-vae_e3m2",
+        revision="da0d2cd7815792fb40d084dbd8ed32b63f153d8d",
+        torch_dtype=Config.DTYPE
+    )
+
+def _load_transformer_model():
+    """Load the transformer model from a specific cached path."""
+    transformer_path = os.path.join(
+        HF_HUB_CACHE,
+        "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a"
+    )
+    return FluxTransformer2DModel.from_pretrained(
+        transformer_path,
+        torch_dtype=Config.DTYPE,
+        use_safetensors=False
+    ).to(memory_format=torch.channels_last)
+
+def _warmup_pipeline(pipeline):
+    """Warm up the pipeline by running it with an empty prompt to initialize internal caches."""
+    for _ in range(3):
+        pipeline(prompt=" ")
+
+def load_pipeline():
+    """
+    Load and configure the diffusion pipeline for text-to-image generation.
+
+    Returns:
+        DiffusionPipeline: The configured pipeline ready for inference.
+    """
+    _clear_gpu_memory()
+
+    # Load individual components
+    text_encoder = _load_text_encoder_model()
+    vae = _load_vae_model()
+    transformer = _load_transformer_model()
+
+    # Assemble the diffusion pipeline
+    pipeline = DiffusionPipeline.from_pretrained(
+        Config.CKPT_ID,
+        vae=vae,
+        revision=Config.CKPT_REVISION,
+        transformer=transformer,
+        text_encoder_2=text_encoder,
+        torch_dtype=Config.DTYPE,
+    ).to(Config.DEVICE)
+
+    # Apply optimizations
+    apply_cache_on_pipe(pipeline)
+    pipeline.to(memory_format=torch.channels_last)
+    pipeline.vae = torch.compile(pipeline.vae, mode="max-autotune")
+    quantize_(pipeline.vae, int8_weight_only())
+    quantize_(pipeline.vae, float8_weight_only())
+
+    # Warm up the pipeline to ensure readiness
+    _warmup_pipeline(pipeline)
+
+    return pipeline
+
+@torch.no_grad()
 def infer(request: TextToImageRequest, pipeline: DiffusionPipeline, generator: Generator) -> Image:
-    """Generate image using the pipeline within a managed context."""
-    context = InferenceContext(pipeline)
-    with context.inference_mode() as p:
-        return p(
-            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"
-        ).images[0]
+    """
+    Generate an image from a text prompt using the diffusion pipeline.
+
+    Args:
+        request (TextToImageRequest): The request containing the prompt and image parameters.
+        pipeline (DiffusionPipeline): The pre-loaded diffusion pipeline.
+        generator (Generator): The random seed generator for reproducibility.
+
+    Returns:
+        Image: The generated image in PIL format.
+    """
+    image = 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"
+    ).images[0]
+    return image
+
+# Initialize environment settings when the module is imported
+_initialize_environment()
+
+# For compatibility with other scripts, alias load_pipeline as load
+load = load_pipeline