generate.py

import torch
import base64
import io
from typing import List, Tuple
from diffusers import StableDiffusionPipeline

# Load model once at startup
MODEL_ID = "runwayml/stable-diffusion-v1-5"
_pipeline = None


def get_pipeline():
    global _pipeline
    if _pipeline is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32
        print(f"Loading Stable Diffusion on {device}...")
        _pipeline = StableDiffusionPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=dtype,
            safety_checker=None
        )
        _pipeline = _pipeline.to(device)
        if device == "cuda":
            _pipeline.enable_attention_slicing()
    return _pipeline


STYLE_MODIFIERS = [
    "warm lighting, cozy atmosphere",
    "bright natural light, airy feel",
    "moody dramatic lighting, luxurious",
    "soft pastel tones, calm and serene",
    "bold accent colors, vibrant energy",
    "neutral earth tones, organic textures",
    "monochrome palette, sleek and clean",
    "golden hour light, warm and inviting"
]


def build_prompt(room_type: str, furniture_style: str, room_theme: str, modifier: str) -> str:
    return (
        f"A beautifully designed {room_theme.lower()} style {room_type.lower()} "
        f"with {furniture_style.lower()} furniture, {modifier}, "
        f"interior design photography, high resolution, 4K, photorealistic, "
        f"professional interior design, architectural digest"
    )


def image_to_base64(image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode("utf-8")


def generate_designs(
    room_type: str,
    furniture_style: str,
    room_theme: str,
    num_variations: int = 5
) -> Tuple[List[str], List[str]]:

    pipeline = get_pipeline()
    images_b64 = []
    prompts_used = []

    modifiers = STYLE_MODIFIERS[:num_variations]

    for modifier in modifiers:
        prompt = build_prompt(room_type, furniture_style, room_theme, modifier)
        negative_prompt = (
            "ugly, blurry, low quality, distorted, deformed, "
            "cluttered, messy, dark, oversaturated, cartoon"
        )

        with torch.no_grad():
            result = pipeline(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=30,
                guidance_scale=7.5,
                width=512,
                height=512
            )

        image = result.images[0]
        images_b64.append(image_to_base64(image))
        prompts_used.append(prompt)

    return images_b64, prompts_used