Spaces:
Running
Running
| 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 | |
| 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> | |
| """ | |
| def index(): | |
| return render_template_string(index_html) | |
| 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) | |