|
from flask import Flask, request, jsonify, send_file |
|
from flask_cors import CORS |
|
from diffusers import StableDiffusionPipeline |
|
import torch |
|
import io |
|
from PIL import Image |
|
|
|
myapp = Flask(__name__) |
|
CORS(myapp) |
|
|
|
|
|
model_url = "https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors" |
|
pipeline = StableDiffusionPipeline.from_single_file(model_url, torch_dtype=torch.float16) |
|
pipeline = pipeline.to("cpu") |
|
|
|
@myapp.route('/') |
|
def home(): |
|
return "Stable Diffusion API is running!" |
|
|
|
@myapp.route('/generate', methods=['POST']) |
|
def generate(): |
|
prompt = request.form.get('prompt') |
|
if not prompt: |
|
return jsonify({"error": "No prompt provided!"}), 400 |
|
|
|
|
|
image = pipeline(prompt).images[0] |
|
|
|
|
|
img_io = io.BytesIO() |
|
image.save(img_io, 'PNG') |
|
img_io.seek(0) |
|
|
|
|
|
return send_file(img_io, mimetype='image/png', as_attachment=True, attachment_filename='generated_image.png') |
|
|
|
if __name__ == '__main__': |
|
myapp.run(host="0.0.0.0", port=8080, debug=True) |