Spaces:
Running
Running
File size: 3,410 Bytes
a2384b5 6cd0d14 435f67e 6cd0d14 435f67e 6cd0d14 435f67e 6cd0d14 435f67e 6cd0d14 435f67e 6cd0d14 435f67e 6cd0d14 700fd72 1fb1361 700fd72 dbf2966 6cd0d14 dbf2966 435f67e a2384b5 6cd0d14 a2384b5 6cd0d14 a2384b5 6cd0d14 a2384b5 6cd0d14 a2384b5 dbf2966 |
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 |
from flask import Flask, request, jsonify, send_file, render_template_string, make_response
import torch
from diffusers import StableDiffusionPipeline
import io
import random
from PIL import Image
from deep_translator import GoogleTranslator
app = Flask(__name__)
# Load the model once globally
model_id = "Ojimi/anime-kawai-diffusion" # Replace with appropriate model if needed
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
# Function to generate the image locally using the Diffusers pipeline
def generate_image_locally(prompt, negative_prompt="", steps=35, cfg_scale=7, seed=-1, width=1024, height=1024):
if not prompt:
return None, "Prompt is required"
# Translate the prompt from Russian to English (if needed)
prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
print(f"Translated prompt: {prompt}")
# Add artistic flair to the prompt
prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
# Set random seed if not provided
generator = torch.manual_seed(seed if seed != -1 else random.randint(1, 1000000000))
try:
# Generate the image using the Diffusers pipeline
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=steps,
guidance_scale=cfg_scale,
generator=generator,
height=height,
width=width,
).images[0]
return image, None
except Exception as e:
return None, f"Error generating image: {str(e)}"
# Content-Security-Policy headers for added security
@app.after_request
def add_security_headers(response):
response.headers['Content-Security-Policy'] = (
"default-src 'self'; "
"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">
<head>
<title>Kawai Diffusion</title>
</head>
<body>
<h1>Welcome to Kawai Diffusion</h1>
<form action="/generate" method="get">
<label for="prompt">Prompt:</label>
<input type="text" id="prompt" name="prompt" required><br><br>
<button type="submit">Generate</button>
</form>
</body>
</html>
"""
@app.route('/')
def index():
return render_template_string(index_html)
@app.route('/generate', methods=['GET'])
def generate():
prompt = request.args.get("prompt", "")
negative_prompt = request.args.get("negative_prompt", "")
steps = int(request.args.get("steps", 35))
cfg_scale = float(request.args.get("cfg_scale", 7))
seed = int(request.args.get("seed", -1))
width = int(request.args.get("width", 1024))
height = int(request.args.get("height", 1024))
# Generate the image locally
image, error = generate_image_locally(prompt, negative_prompt, steps, cfg_scale, seed, width, height)
if error:
return jsonify({"error": error}), 400
# Send the generated image as a PNG file
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)
|