File size: 3,088 Bytes
af7056a 848b0e8 004975c ccd0584 004975c af7056a 848b0e8 af7056a 227771d ccd0584 af7056a ccd0584 af7056a ccd0584 df466fb ccd0584 bbabb76 af7056a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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)
|