completed text to image pipeline using SD

app/generator.py

@@ -1 +1,83 @@
-"""Auto-generated placeholder module for Stable Diffusion Image Generator."""
+"""Image generation wrapper around a loaded StableDiffusionPipeline.
+
+Provides:
+- generate_image(...) -> (PIL.Image, metadata)
+- deterministic seed handling
+"""
+
+import time
+from typing import Any, Dict, Optional
+
+import torch
+
+from app.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def _validate_resolution(width: int, height: int):
+    # clamp and snap to multiples of 64 (SD requirement)
+    width = max(256, min(width, 768))
+    height = max(256, min(height, 768))
+    width = (width // 64) * 64
+    height = (height // 64) * 64
+    return int(width), int(height)
+
+
+def generate_image(
+    pipe,
+    prompt: str,
+    negative_prompt: Optional[str] = None,
+    steps: int = 30,
+    guidance_scale: float = 7.5,
+    width: int = 512,
+    height: int = 512,
+    seed: Optional[int] = None,
+    device: str = "cuda",
+):
+    """Generate a single image and return (PIL.Image, metadata dict)."""
+    start = time.time()
+    width, height = _validate_resolution(width, height)
+
+    # Generator for reproducibility
+    if seed is None:
+        # create a new seed and use it
+        seed = int(torch.seed() & ((1 << 63) - 1))
+    gen = torch.Generator(device if device != "cpu" else "cpu").manual_seed(int(seed))
+
+    logger.info(
+        (
+            f"Generating: steps={steps}, cfg={guidance_scale},\
+        res={width}x{height}, seed={seed}"
+        )
+    )
+
+    # Use autocast for speed/precision management
+    device_type = "cuda" if device != "cpu" else "cpu"
+    with torch.autocast(device_type=device_type):
+        result = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance_scale),
+            width=width,
+            height=height,
+            generator=gen,
+        )
+
+    img = result.images[0]  # PIL image
+    elapsed = time.time() - start
+
+    metadata: Dict[str, Any] = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "steps": steps,
+        "guidance_scale": guidance_scale,
+        "width": width,
+        "height": height,
+        "seed": int(seed),
+        "elapsed_seconds": elapsed,
+    }
+
+    logger.info(f"Generation finished in {elapsed:.2f}s")
+    return img, metadata

app/pipeline.py

@@ -1 +1,122 @@
-"""Auto-generated placeholder module for Stable Diffusion Image Generator."""
+"""Model pipeline loader for Stable Diffusion (HuggingFace Diffusers).
+
+load_pipeline(...) returns a GPU-ready pipeline with memory optimizations.
+"""
+
+import os
+from typing import Optional
+
+import torch
+from diffusers import (
+    DPMSolverMultistepScheduler,
+    StableDiffusionPipeline,
+)
+from dotenv import load_dotenv
+
+from app.utils.logger import get_logger
+
+logger = get_logger(__name__)
+load_dotenv()
+
+
+def _try_enable_xformers(pipe):
+    try:
+        if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
+            pipe.enable_xformers_memory_efficient_attention()
+            logger.info("Enabled xFormers memory-efficient attention.")
+        else:
+            logger.info("xFormers not available via API; skipping.")
+    except Exception as err:
+        logger.info(f"xFormers not enabled: {err}")
+
+
+def load_pipeline(
+    model_id: str = "runwayml/stable-diffusion-v1-5",
+    device: str = "cuda",
+    use_fp16: bool = True,
+    enable_xformers: bool = False,
+    torch_dtype: Optional[torch.dtype] = None,
+    scheduler=None,
+):
+    """Load and return an optimized StableDiffusionPipeline."""
+    if torch_dtype is None:
+        torch_dtype = torch.float16 if use_fp16 and device == "cuda" else torch.float32
+
+    if scheduler is None:
+        try:
+            scheduler = DPMSolverMultistepScheduler.from_pretrained(
+                model_id,
+                subfolder="scheduler",
+            )
+        except Exception:
+            scheduler = None
+
+    logger.info(f"Loading pipeline {model_id} " f"dtype={torch_dtype} on {device} ...")
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch_dtype,
+        safety_checker=None,
+        scheduler=scheduler,
+        use_auth_token=os.getenv("HUGGINGFACE_HUB_TOKEN"),
+    )
+
+    pipe = pipe.to(device)
+
+    try:
+        pipe.enable_attention_slicing()
+        logger.info("Enabled attention slicing.")
+    except Exception:
+        logger.info("Attention slicing not available.")
+
+    if enable_xformers:
+        _try_enable_xformers(pipe)
+
+    try:
+        if hasattr(pipe.vae, "enable_tiling"):
+            pipe.vae.enable_tiling()
+            logger.info("Enabled VAE tiling.")
+    except Exception:
+        pass
+
+    if device == "cuda":
+        torch.backends.cudnn.benchmark = True
+
+    logger.info("Pipeline loaded.")
+    return pipe
+
+
+def warmup_pipeline(
+    pipe,
+    prompt: str = "A photo of a cat",
+    height: int = 512,
+    width: int = 512,
+):
+    """Run a quick inference to allocate CUDA kernels and memory."""
+    try:
+        if hasattr(pipe, "parameters"):
+            device = next(pipe.parameters()).device
+        else:
+            device = "cuda"
+
+    except Exception:
+        device = "cuda"
+
+    try:
+        gen = torch.Generator(device if device != "cpu" else "cpu").manual_seed(0)
+
+        logger.info("Warmup: running one-step inference to initialize kernels.")
+
+        _ = pipe(
+            prompt=prompt,
+            num_inference_steps=1,
+            guidance_scale=1.0,
+            height=height,
+            width=width,
+            generator=gen,
+        )
+
+        torch.cuda.empty_cache()
+        logger.info("Warmup complete.")
+    except Exception as err:
+        logger.warning(f"Warmup failed: {err}")

app/utils/logger.py

@@ -1 +1,51 @@
-"""Auto-generated placeholder module for Stable Diffusion Image Generator."""
+"""Centralized logging utility for the project.
+
+Features:
+- Colored console logs
+- File logs (logs/app.log)
+- Timestamped + module-aware output
+"""
+
+import logging
+import os
+from logging.handlers import RotatingFileHandler
+
+LOG_DIR = "logs"
+LOG_FILE = os.path.join(LOG_DIR, "app.log")
+
+os.makedirs(LOG_DIR, exist_ok=True)
+
+
+def get_logger(name: str = "app", level=logging.INFO) -> logging.Logger:
+    """Returns a configured logger instance.
+
+    Safe to call from any module.
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    if logger.hasHandlers():
+        return logger
+
+    # Console handler
+    console_handler = logging.StreamHandler()
+    console_format = (
+        "\033[36m[%(asctime)s] [%(name)s] \
+        [%(levelname)s]\033[0m "
+        "%(message)s"
+    )
+    console_handler.setFormatter(logging.Formatter(console_format, "%Y-%m-%d %H:%M:%S"))
+
+    # File handler
+    file_handler = RotatingFileHandler(
+        LOG_FILE,
+        maxBytes=5_000_000,
+        backupCount=3,
+    )
+    file_format = "[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s"
+    file_handler.setFormatter(logging.Formatter(file_format, "%Y-%m-%d %H:%M:%S"))
+
+    logger.addHandler(console_handler)
+    logger.addHandler(file_handler)
+
+    return logger

requirements.txt

@@ -8,6 +8,7 @@ torchaudio==2.5.1
 
 
 # HUGGINGFACE DIFFUSION ECOSYSTEM
+huggingface_hub==0.20.3
 diffusers==0.26.3
 transformers==4.39.3
 accelerate==0.28.0