image_generation.py

import torch
from diffusers import StableDiffusionPipeline
import os

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe = pipe.to("cpu")

def generate_images(topic, content_format, tone, num_images=3):
    os.makedirs("images", exist_ok=True)
    prompt = f"{tone} {content_format} illustration about {topic}, colorful modern design"
    images = []
    for i in range(num_images):
        image = pipe(prompt).images[0]
        img_path = f"images/img_{i+1}.png"
        image.save(img_path)
        images.append(img_path)
    return images