File size: 4,091 Bytes
a2384b5
ef72775
6a3ef3a
ef72775
 
 
435f67e
 
 
ef72775
 
435f67e
ef72775
 
 
 
435f67e
ef72775
95d6160
ef72775
 
 
435f67e
ef72775
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1fb1361
700fd72
 
dbf2966
 
ef72775
dbf2966
435f67e
 
 
 
 
 
a2384b5
95d6160
ef72775
0279fde
ef72775
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2384b5
 
ef72775
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
from flask import Flask, request, jsonify, send_file, render_template_string, make_response
import requests
import io
import random
from PIL import Image
from deep_translator import GoogleTranslator

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/Ojimi/anime-kawai-diffusion"
timeout = 3000  # タイムアウトを300秒に設定

# Function to query the API and return the generated image
def query(prompt, negative_prompt="", steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
    if not prompt:
        return None, "Prompt is required"

    key = random.randint(0, 999)
    
    # Translate the prompt from Russian to English if necessary
    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
    print(f'Generation {key} translation: {prompt}')

    # Add some extra flair to the prompt
    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'Generation {key}: {prompt}')
    
    payload = {
        "inputs": prompt,
        "is_negative": False,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "parameters": {
            "width": width,
            "height": height
        }
    }

    for attempt in range(3):  # 最大3回の再試行
        try:
            # Authorization header is removed
            response = requests.post(API_URL, json=payload, timeout=timeout)
            if response.status_code != 200:
                return None, f"Error: Failed to get image. Status code: {response.status_code}, Details: {response.text}"

            image_bytes = response.content
            image = Image.open(io.BytesIO(image_bytes))
            return image, None
        except requests.exceptions.Timeout:
            if attempt < 2:  # 最後の試行でない場合は再試行
                print("Timeout occurred, retrying...")
                continue
            return None, "Error: The request timed out. Please try again."
        except requests.exceptions.RequestException as e:
            return None, f"Request Exception: {str(e)}"
        except Exception as e:
            return None, f"Error when trying to open the image: {e}"

# Content-Security-Policyヘッダーを設定するための関数
@app.after_request
def add_security_headers(response):
    response.headers['Content-Security-Policy'] = (
        "default-src 'self'; "
        "connect-src 'self' ^https?:\/\/[\w.-]+\.[\w.-]+(\/[\w.-]*)*(\?[^\s]*)?$"
        "img-src 'self' data:; "
        "style-src 'self' 'unsafe-inline'; "
        "script-src 'self' 'unsafe-inline'; "
    )
    return response

# HTML template for the index page
index_html = """
<!DOCTYPE html>
<html lang="ja">
kawai diffusion
</html>
"""

@app.route('/')
def index():
    return render_template_string(index_html)

@app.route('/generate', methods=['GET'])
def generate_image():
    prompt = request.args.get("prompt", "")
    negative_prompt = request.args.get("negative_prompt", "bad hands, low quality, low quality, unnatural, dirty eyes, distorted eyes, low quality eyes, distorted arms, distorted legs, distorted face, 3 legs, 3 hands, anatomically Bad, 6 fingers, 7 fingers, less fingers, more fingers,")
    steps = int(request.args.get("steps", 35))
    cfg_scale = float(request.args.get("cfgs", 7))
    sampler = request.args.get("sampler", "DPM++ 2M Karras")
    strength = float(request.args.get("strength", 0.7))
    seed = int(request.args.get("seed", -1))
    width = int(request.args.get("width", 1024))
    height = int(request.args.get("height", 1024))

    image, error = query(prompt, negative_prompt, steps, cfg_scale, sampler, seed, strength, width, height)

    if error:
        return jsonify({"error": error}), 400

    img_bytes = io.BytesIO()
    image.save(img_bytes, format='PNG')
    img_bytes.seek(0)
    return send_file(img_bytes, mimetype='image/png')

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