app.py

from flask import Flask, request, jsonify
import torch
from transformers import (
    UNet2DConditionModel, 
    AutoTokenizer, 
    CLIPTextModel, 
    CLIPTextModelWithProjection, 
    CLIPVisionModelWithProjection
)
from PIL import Image
import base64
from io import BytesIO

app = Flask(__name__)

# Global variables for models to load them once at startup
unet = None
tokenizer_one = None
tokenizer_two = None
noise_scheduler = None
text_encoder_one = None
text_encoder_two = None
image_encoder = None
vae = None
UNet_Encoder = None

# Load models once at startup
def load_models():
    global unet, tokenizer_one, tokenizer_two, noise_scheduler
    global text_encoder_one, text_encoder_two, image_encoder, vae, UNet_Encoder

    if unet is None:
        # Load models only when required to reduce memory usage
        unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-v1-4")
    
    if tokenizer_one is None:
        tokenizer_one = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14")
    
    if tokenizer_two is None:
        tokenizer_two = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14-336")
    
    if noise_scheduler is None:
        noise_scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4")

    if text_encoder_one is None:
        text_encoder_one = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

    if text_encoder_two is None:
        text_encoder_two = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-large-patch14-336")
    
    if image_encoder is None:
        image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14")

    if vae is None:
        vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-v1-4")

    if UNet_Encoder is None:
        UNet_Encoder = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-v1-4")

# Helper function to process base64 image
def decode_image(image_base64):
    image_data = base64.b64decode(image_base64)
    image = Image.open(BytesIO(image_data)).convert("RGB")
    return image

# Helper function to encode image to base64
def encode_image(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

# Route for image processing
@app.route('/process_image', methods=['POST'])
def process_image():
    data = request.json

    # Load the models (this will only happen once)
    load_models()

    # Extract the image from the request
    image_base64 = data.get('image_base64')
    if not image_base64:
        return jsonify({"error": "No image provided"}), 400

    image = decode_image(image_base64)

    # Perform inference with the models (example, modify as needed)
    processed_image = image  # Placeholder for actual image processing

    # Return the processed image as base64
    processed_image_base64 = encode_image(processed_image)
    return jsonify({"processed_image": processed_image_base64})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)