File size: 2,469 Bytes
faae644
a650a4e
 
 
 
 
faae644
82dd938
 
a650a4e
 
 
91e0236
a650a4e
 
 
faae644
 
a650a4e
 
 
 
 
 
 
 
 
5065d5c
 
 
a650a4e
 
 
82dd938
 
 
faae644
82dd938
 
 
 
 
 
 
 
cc1489b
 
82dd938
 
 
cc1489b
a650a4e
 
 
 
 
 
 
 
 
faae644
cc1489b
a650a4e
 
cc1489b
 
a650a4e
cc1489b
a650a4e
 
 
 
 
 
 
 
faae644
 
 
 
a650a4e
faae644
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
from flask import Flask, request, jsonify, send_from_directory
import requests
import io
import random
import os
from PIL import Image
from datetime import datetime
import string
import re

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/openskyml/midjourney-mini"
API_TOKEN = os.getenv("HF_READ_TOKEN")  # it is free
headers = {"Authorization": f"Bearer {API_TOKEN}"}

TEMP_DIR = "temp"

def query(prompt, is_negative=False, steps=1, cfg_scale=6, seed=None):
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed is not None else random.randint(-1, 2147483647)
    }

    response = requests.post(API_URL, headers=headers, json=payload)
    response.raise_for_status()
    image_bytes = response.content
    image = Image.open(io.BytesIO(image_bytes))
    return image

def sanitize_filename(filename):
    # Replace spaces with underscores
    filename = filename.replace(" ", "_")

    # Remove special characters
    sanitized_filename = re.sub(r'[^\w\s-]', '', filename)
    sanitized_filename = re.sub(r'[-\s]+', '-', sanitized_filename)

    return sanitized_filename

def save_image(image, prompt):
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    sanitized_prompt = sanitize_filename(prompt)
    filename = f"{sanitized_prompt}-{random.randint(1, 100000)}-{timestamp}.png"
    filepath = os.path.join(TEMP_DIR, filename)
    image.save(filepath, format='PNG')
    return filename

@app.route("/generate", methods=["POST"])
def generate():
    try:
        data = request.get_json()
        prompt = data["prompt"]
        negative_prompt = data.get("negative_prompt", "")
        is_negative = True if negative_prompt else False

        image = query(prompt, is_negative=is_negative)
        filename = save_image(image, prompt)

        response = {
            "success": True,
            "image": f"https://mrdonstuff-dalle-3-xl-api.hf.space/temp/{filename}",
            "filename": filename
        }

    except Exception as e:
        response = {
            "success": False,
            "error": str(e)
        }

    return jsonify(response)

@app.route("/temp/<filename>")
def show_image(filename):
    return send_from_directory(TEMP_DIR, filename)

if __name__ == "__main__":
    if not os.path.exists(TEMP_DIR):
        os.makedirs(TEMP_DIR)

    app.run(debug=True, host="0.0.0.0", port=7860)