img2img-turbo / draw_server.py
Inmental's picture
Upload folder using huggingface_hub
343e5a8 verified
from flask import Flask, render_template_string, request, send_file, jsonify
from PIL import Image
from io import BytesIO
import os
import threading
import gradio as gr
app = Flask(__name__)
# File paths
SKETCH_PATH = "sketch.png"
OUTPUT_PATH = "output.png"
# Route to upload the sketch
@app.route('/upload_sketch', methods=['POST'])
def upload_sketch():
sketch = request.files['sketch']
sketch.save(SKETCH_PATH)
return jsonify({"status": "success", "message": "Sketch uploaded successfully."})
# Route to upload the output image
@app.route('/upload_output', methods=['POST'])
def upload_output():
output = request.files['output']
output.save(OUTPUT_PATH)
return jsonify({"status": "success", "message": "Output uploaded successfully."})
# Route to get the latest sketch
@app.route('/get_sketch', methods=['GET'])
def get_sketch():
if os.path.exists(SKETCH_PATH):
return send_file(SKETCH_PATH, mimetype='image/png')
return jsonify({"status": "error", "message": "Sketch not found."}), 404
# Route to get the latest output image
@app.route('/get_output', methods=['GET'])
def get_output():
if os.path.exists(OUTPUT_PATH):
return send_file(OUTPUT_PATH, mimetype='image/png')
return jsonify({"status": "error", "message": "Output not found."}), 404
# Route to render the preview page
@app.route('/')
def index():
# HTML template for the preview page
html_template = """
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Preview Page</title>
<style>
body, html {
margin: 0;
padding: 0;
height: 100%;
background-color: black;
}
.full-screen-image {
width: 100%;
height: 100%;
object-fit: contain;
}
</style>
<script>
function refreshImage() {
var img = document.getElementById("output-image");
img.src = "/get_output?" + new Date().getTime();
}
// Auto-refresh every 2 seconds to show the latest image
setInterval(refreshImage, 2000);
</script>
</head>
<body>
<img id="output-image" src="/get_output" class="full-screen-image">
</body>
</html>
"""
return render_template_string(html_template)
# Function to run the Gradio drawing interface
def start_gradio_interface():
# Setup the Gradio interface for drawing
with gr.Blocks(css="style.css") as demo:
gr.Markdown("# Drawing Page")
# Hidden buttons for canvas actions
line = gr.Checkbox(label="line", value=False, elem_id="cb-line")
eraser = gr.Checkbox(label="eraser", value=False, elem_id="cb-eraser")
with gr.Row(elem_id="main_row"):
with gr.Column(elem_id="column_input"):
gr.Markdown("## INPUT", elem_id="input_header")
image = gr.Image(
source="canvas",
tool="color-sketch",
type="pil",
image_mode="L",
invert_colors=True,
shape=(512, 512),
brush_radius=4,
height=512,
width=512,
brush_color="#000000",
interactive=True,
show_download_button=True,
elem_id="input_image",
show_label=False,
)
upload_btn = gr.UploadButton(label="Upload your sketch", file_types=["image"], file_count="single")
def upload_sketch(image):
buffered = BytesIO()
image.save(buffered, format="PNG")
buffered.seek(0)
response = requests.post(f"http://localhost:5000/upload_sketch", files={"sketch": buffered})
return response.json()
upload_btn.click(upload_sketch, inputs=image, outputs=None)
gr.HTML(
"""
<div class="button-row">
<div id="my-div-pencil" class="pad2"> <button id="my-toggle-pencil" onclick="return togglePencil(this)"></button> </div>
<div id="my-div-eraser" class="pad2"> <button id="my-toggle-eraser" onclick="return toggleEraser(this)"></button> </div>
<div class="pad2"> <button id="my-button-undo" onclick="return UNDO_SKETCH_FUNCTION(this)"></button> </div>
<div class="pad2"> <button id="my-button-clear" onclick="return DELETE_SKETCH_FUNCTION(this)"></button> </div>
<div class="pad2"> <button href="TODO" download="image" id="my-button-down" onclick='return theSketchDownloadFunction()'></button> </div>
</div>
"""
)
run_button = gr.Button("Run", max_width=20)
prompt = gr.Textbox(label="Prompt", value="", show_label=True)
demo.launch(server_name="0.0.0.0", server_port=7860)
# Run Flask and Gradio in parallel
if __name__ == "__main__":
threading.Thread(target=start_gradio_interface).start()
app.run(host='0.0.0.0', port=5000)