Spaces:
Running
Running
| 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ヘッダーを設定するための関数 | |
| 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> | |
| """ | |
| def index(): | |
| return render_template_string(index_html) | |
| 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) |