imagegen_api.py

import os
from flask import jsonify, send_file, request
from io import BytesIO
from PIL import Image
from main import *
import torch, base64

def generate_image(prompt, output_path="output_image.png"):
    if imagegen_model is None: return {"error": "Image generation model not initialized."}
    generator = torch.Generator(device=device).manual_seed(0)
    with torch.no_grad(): image = imagegen_model(prompt, generator=generator,).images[0]; image.save(output_path); return output_path

def imagegen_api(prompt):
    output_file = generate_image(prompt)
    if isinstance(output_file, dict) and "error" in output_file: return {"error": output_file["error"]}
    image_io = BytesIO(); pil_image = Image.open(output_file); pil_image.save(image_io, 'PNG'); image_base64 = base64.b64encode(image_io.getvalue()).decode('utf-8')
    os.remove(output_file); return {"image_base64": image_base64, "mimetype": "image/png"}