59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
from flask import Flask, request, jsonify
|
|
from diffusers import DDPMScheduler, DiffusionPipeline
|
|
import torch
|
|
import base64
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
|
|
app = Flask(__name__)
|
|
|
|
scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")
|
|
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", scheduler=scheduler)
|
|
|
|
print("Loading UNet weights...")
|
|
ckpt_path = './models/model.bin'
|
|
try:
|
|
unet_state_dict = torch.load(ckpt_path)
|
|
pipe.unet.load_state_dict(unet_state_dict)
|
|
if torch.cuda.is_available():
|
|
pipe = pipe.to("cuda")
|
|
pipe.to(torch.float16)
|
|
print("UNet weights loaded successfully and model moved to CUDA.")
|
|
else:
|
|
print("CUDA is not available. Running on CPU (will be slow).")
|
|
pipe = pipe.to("cpu")
|
|
pipe.to(torch.float32)
|
|
except FileNotFoundError:
|
|
print(f"Error: Model not found at {ckpt_path}. Please ensure the file exists.")
|
|
exit()
|
|
except Exception as e:
|
|
print(f"An error occurred while loading UNet weights: {e}")
|
|
exit()
|
|
|
|
@app.route('/api/generate', methods=['POST'])
|
|
def generate_image():
|
|
prompt = request.json.get('prompt')
|
|
if not prompt:
|
|
return jsonify({"error": "Prompt is required in the request body."}), 400
|
|
|
|
print(f"Generating image for prompt: '{prompt}'...")
|
|
try:
|
|
image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[999]).images[0]
|
|
image = image.resize((128, 128), Image.LANCZOS)
|
|
|
|
# Convert image to base64
|
|
buffered = BytesIO()
|
|
image.save(buffered, format="JPEG") # You can choose JPEG for smaller size, but PNG is lossless
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
return jsonify({"image": img_str})
|
|
|
|
except Exception as e:
|
|
print(f"Error during image generation: {e}")
|
|
return jsonify({"error": "An error occurred during image generation."}), 500
|
|
|
|
if __name__ == '__main__':
|
|
# Make sure to run this with 'python app.py'
|
|
# It will be accessible at http://127.0.0.1:5000/
|
|
app.run(host='127.0.0.1', port=5000, debug=True) # debug=True is good for development
|