You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.5 KiB
78 lines
2.5 KiB
import json
|
|
import traceback
|
|
from flask import Flask, render_template, request, jsonify, send_from_directory, Response
|
|
from sd_pipeline import pipeline
|
|
from models import GenerationOptions
|
|
|
|
app = Flask(__name__)
|
|
|
|
|
|
@app.route("/")
|
|
def index():
|
|
return render_template("index.html")
|
|
|
|
|
|
@app.route("/generate", methods=["POST"])
|
|
def generate():
|
|
data = request.get_json()
|
|
|
|
prompt = data.get("prompt", "")
|
|
if not prompt:
|
|
return jsonify({"success": False, "error": "Prompt is required"}), 400
|
|
|
|
# Parse and validate seed
|
|
seed = data.get("seed")
|
|
if seed is not None and seed != "":
|
|
seed = int(seed)
|
|
else:
|
|
seed = None
|
|
|
|
# Parse and clamp numeric values
|
|
width = data.get("width")
|
|
height = data.get("height")
|
|
if width:
|
|
width = max(256, min(2048, int(width) // 8 * 8)) # must be multiple of 8
|
|
if height:
|
|
height = max(256, min(2048, int(height) // 8 * 8))
|
|
|
|
options = GenerationOptions(
|
|
prompt=prompt,
|
|
negative_prompt=data.get("negative_prompt", ""),
|
|
seed=seed,
|
|
steps=max(1, min(100, int(data.get("steps", 20)))),
|
|
guidance_scale=max(1.0, min(20.0, float(data.get("guidance_scale", 7.5)))),
|
|
count=max(1, min(10, int(data.get("count", 1)))),
|
|
add_quality_keywords=data.get("add_quality_keywords", True),
|
|
increment_seed=data.get("increment_seed", True),
|
|
vary_guidance=data.get("vary_guidance", False),
|
|
guidance_low=max(1.0, min(20.0, float(data.get("guidance_low", 5.0)))),
|
|
guidance_high=max(1.0, min(20.0, float(data.get("guidance_high", 12.0)))),
|
|
vary_steps=data.get("vary_steps", False),
|
|
steps_low=max(1, min(100, int(data.get("steps_low", 20)))),
|
|
steps_high=max(1, min(100, int(data.get("steps_high", 80)))),
|
|
width=width,
|
|
height=height,
|
|
)
|
|
|
|
def generate_events():
|
|
try:
|
|
for result in pipeline.generate_stream(options):
|
|
yield f"data: {json.dumps(result.to_dict())}\n\n"
|
|
yield f"data: {json.dumps({'done': True})}\n\n"
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
|
|
|
return Response(generate_events(), mimetype='text/event-stream')
|
|
|
|
|
|
@app.route("/out/<path:filename>")
|
|
def serve_image(filename):
|
|
return send_from_directory("out", filename)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Loading model on startup...")
|
|
pipeline.load()
|
|
print("Starting web server...")
|
|
app.run(host="127.0.0.1", port=5000, debug=False, threaded=True)
|
|
|