Spaces:
Build error
Build error
| """EuropaLex Image Generation Engine — diffusers Flux2KleinPipeline.""" | |
| from __future__ import annotations | |
| import logging | |
| from pathlib import Path | |
| import torch | |
| from core.types import ImageResult | |
| logger = logging.getLogger(__name__) | |
| class ImageGenEngine: | |
| """Image generation using diffusers Flux2KleinPipeline. | |
| Lazy-loads the pipeline on first generation call, unloads after completion. | |
| Only one instance can be active at a time (enforced by EnginePool). | |
| """ | |
| def __init__(self, device: str = "cuda"): | |
| """Initialize the image engine. | |
| Args: | |
| device: 'cuda', 'mps', or 'cpu'. | |
| """ | |
| self.device = device | |
| self._pipeline = None | |
| self._loaded = False | |
| def _load_pipeline(self) -> None: | |
| """Lazy-load the Flux2Klein pipeline from HF Hub (cached locally).""" | |
| if self._loaded: | |
| return | |
| try: | |
| from diffusers import Flux2KleinPipeline | |
| except ImportError: | |
| raise ImportError( | |
| "diffusers package not installed. Run: pip install diffusers" | |
| ) | |
| torch_dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 | |
| logger.info("Loading Flux2Klein from HF Hub (cached in ~/.cache/huggingface/)") | |
| self._pipeline = Flux2KleinPipeline.from_pretrained( | |
| "black-forest-labs/FLUX.2-klein-4B", | |
| torch_dtype=torch_dtype, | |
| ) | |
| self._pipeline.enable_model_cpu_offload() | |
| self._loaded = True | |
| logger.info("Flux2Klein pipeline loaded on %s", self.device) | |
| def generate(self, prompts: list[str], output_dir: Path) -> ImageResult: | |
| """Generate images for a batch of prompts. | |
| Args: | |
| prompts: List of text prompts for image generation. | |
| output_dir: Directory to save .png files. | |
| Returns: | |
| ImageResult with absolute paths to generated image files. | |
| """ | |
| self._load_pipeline() | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| image_paths = [] | |
| for i, prompt in enumerate(prompts): | |
| try: | |
| images = self._pipeline( | |
| prompt=prompt, | |
| num_inference_steps=4, | |
| guidance_scale=1.0, | |
| width=240, | |
| height=160, | |
| ) | |
| if images.images and len(images.images) > 0: | |
| img_path = output_dir / f"image_{i}.png" | |
| images.images[0].save(str(img_path)) | |
| image_paths.append(str(img_path.resolve())) | |
| logger.debug("Saved image to %s", img_path) | |
| else: | |
| logger.warning("Empty image output for prompt: %s", prompt[:50]) | |
| image_paths.append(None) | |
| except Exception as e: | |
| logger.error("Image generation failed for prompt '%s': %s", prompt[:50], e) | |
| image_paths.append(None) | |
| return ImageResult(image_paths=list(image_paths)) | |
| def unload(self) -> None: | |
| """Unload the pipeline and free GPU memory.""" | |
| if self._pipeline is not None: | |
| del self._pipeline | |
| self._pipeline = None | |
| self._loaded = False | |
| try: | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| logger.info("Flux2Klein pipeline unloaded") | |