draw.py

import random
import numpy as np
from PIL import Image, ImageOps
import base64
from io import BytesIO
import torch
import torchvision.transforms.functional as F
import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration
from flask import Flask, request, jsonify, render_template_string, send_file
from flask_cors import CORS
import threading
import hashlib
import signal
import sys
import os

# Load models
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cuda")

# Pix2Pix model placeholder (Assume you have this model correctly implemented)
class Pix2Pix_Turbo:
    def __init__(self, mode):
        pass

    def __call__(self, c_t, prompt, deterministic, r, noise_map):
        # Dummy image processing function for demonstration purposes
        return c_t

pix2pix_model = Pix2Pix_Turbo("sketch_to_image_stochastic")

# Flask application setup
app = Flask(__name__)
CORS(app)  # Handle CORS issues

# Global Constants and Configuration
STYLE_LIST = [
    {"name": "Cinematic", "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy"},
    {"name": "3D Model", "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting"},
    {"name": "Anime", "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed"},
    {"name": "Digital Art", "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed"},
    {"name": "Photographic", "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed"},
    {"name": "Pixel art", "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics"},
    {"name": "Fantasy art", "prompt": "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy"},
    {"name": "Neonpunk", "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional"},
    {"name": "Manga", "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style"},
]

STYLES = {style["name"]: style["prompt"] for style in STYLE_LIST}
STYLE_NAMES = list(STYLES.keys())
DEFAULT_STYLE_NAME = "Fantasy art"
MAX_SEED = np.iinfo(np.int32).max

# Paths for storing sketches and outputs
SKETCH_PATH = "sketch.png"
OUTPUT_PATH = "output.png"

# Image processing function
def run(image, prompt, prompt_template, style_name, seed, val_r):
    if not prompt.strip():
        prompt = "Generated by drawing tool"
    
    prompt = prompt_template.replace("{prompt}", prompt)
    image = image.convert("RGB")
    image_tensor = F.to_tensor(image) > 0.5

    with torch.no_grad():
        c_t = image_tensor.unsqueeze(0).to("cuda").float()
        torch.manual_seed(seed)
        noise = torch.randn((1, 4, c_t.shape[2] // 8, c_t.shape[3] // 8), device=c_t.device)
        output_image = pix2pix_model(c_t, prompt, deterministic=False, r=val_r, noise_map=noise)

    output_pil = F.to_pil_image(output_image[0].cpu() * 0.5 + 0.5)
    output_pil.save(OUTPUT_PATH)  # Save the output image

    buffered = BytesIO()
    output_pil.save(buffered, format="PNG")
    output_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return output_data

# Flask route to handle image processing
@app.route('/process-image', methods=['POST'])
def process_image():
    try:
        data = request.get_json()
        image_data = data.get("image", "").split(",")[1]
        image = Image.open(BytesIO(base64.b64decode(image_data))).convert("RGB")

        # Process the image
        output_image_uri = run(
            image,
            data.get("prompt", ""),
            STYLES.get(data.get("style_name", DEFAULT_STYLE_NAME)),
            data.get("style_name", DEFAULT_STYLE_NAME),
            int(data.get("seed", random.randint(0, MAX_SEED))),
            float(data.get("val_r", 0.4))
        )
        return jsonify({"image": output_image_uri})

    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Flask route to serve the sketch image
@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

# Flask route to serve the 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

# HTML page for drawing
@app.route('/')
def draw_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>Drawing Page</title>

        <style>

            body, html {

                margin: 0;

                padding: 0;

                height: 100%;

                display: flex;

                justify-content: center;

                align-items: center;

                background-color: #f0f0f0;

            }

            .canvas-container {

                border: 2px solid black;

                position: relative;

            }

            .toolbar {

                display: flex;

                justify-content: center;

                margin-bottom: 10px;

            }

            button {

                margin-right: 5px;

            }

            canvas {

                cursor: crosshair;

            }

        </style>

    </head>

    <body>

        <div class="toolbar">

            <button id="brush" onclick="setTool('brush')">Brush</button>

            <button id="line" onclick="setTool('line')">Line</button>

            <button id="eraser" onclick="setTool('eraser')">Eraser</button>

            <button id="clear" onclick="clearCanvas()">Clear</button>

            <input type="color" id="colorPicker" value="#000000">

            <input type="range" id="brushSize" min="1" max="20" value="4">

        </div>

        <div class="canvas-container">

            <canvas id="drawingCanvas" width="800" height="600"></canvas>

        </div>

        <script>

            let canvas = document.getElementById('drawingCanvas');

            let ctx = canvas.getContext('2d');

            let drawing = false;

            let tool = 'brush';

            let lastX = 0, lastY = 0;



            canvas.addEventListener('mousedown', (e) => {

                drawing = true;

                [lastX, lastY] = [e.offsetX, e.offsetY];

            });



            canvas.addEventListener('mousemove', draw);

            canvas.addEventListener('mouseup', () => {

                drawing = false;

                sendDrawingToBackend();

            });

            canvas.addEventListener('mouseout', () => drawing = false);



            function draw(e) {

                if (!drawing) return;



                ctx.strokeStyle = document.getElementById('colorPicker').value;

                ctx.lineWidth = document.getElementById('brushSize').value;

                ctx.lineJoin = 'round';

                ctx.lineCap = 'round';



                ctx.beginPath();

                ctx.moveTo(lastX, lastY);

                ctx.lineTo(e.offsetX, e.offsetY);

                ctx.stroke();

                [lastX, lastY] = [e.offsetX, e.offsetY];

            }



            function setTool(selectedTool) {

                tool = selectedTool;

                ctx.globalCompositeOperation = (tool === 'eraser') ? 'destination-out' : 'source-over';

            }



            function clearCanvas() {

                ctx.clearRect(0, 0, canvas.width, canvas.height);

            }



            function sendDrawingToBackend() {

                let dataURL = canvas.toDataURL('image/png');

                fetch('/process-image', {

                    method: 'POST',

                    headers: {

                        'Content-Type': 'application/json',

                    },

                    body: JSON.stringify({ image: dataURL }),

                })

                .then(response => response.json())

                .then(data => console.log('Image processed', data))

                .catch(error => console.error('Error processing image:', error));

            }

        </script>

    </body>

    </html>

    """
    return render_template_string(html_template)

# HTML page for previewing the processed image
@app.route('/preview')
def 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)

def signal_handler(sig, frame):
    print("Ctrl+C pressed, shutting down.")
    sys.exit(0)

# Register the signal handler for Ctrl+C
signal.signal(signal.SIGINT, signal_handler)

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=2073)