Spaces:
Runtime error
Runtime error
| from flask import Flask, request, jsonify, send_file, render_template_string, make_response | |
| import torch | |
| from diffusers import StableDiffusionPipeline | |
| from PIL import Image | |
| import io | |
| import random | |
| from huggingface_hub import login | |
| import os | |
| login(token= os.getenv("tk")) | |
| app = Flask(__name__) | |
| MODEL_ID = "deadman44/SD_Anime_Merged_Models" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Diffusersモデルの読み込み | |
| print("Loading the model...") | |
| pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16, safety_checker=None) | |
| pipe = pipe.to(DEVICE) | |
| print("Model loaded successfully.") | |
| # 画像生成関数 | |
| def generate_image_with_diffusers(prompt, negative_prompt="", steps=35, cfg_scale=7, seed=-1, width=512, height=512): | |
| if not prompt: | |
| return None, "Prompt is required" | |
| # シードを固定（-1の場合はランダムに生成） | |
| generator = torch.manual_seed(seed if seed != -1 else random.randint(0, 1000000000)) | |
| # 画像生成 | |
| try: | |
| result = pipe( | |
| prompt, | |
| negative_prompt=negative_prompt, | |
| num_inference_steps=steps, | |
| guidance_scale=cfg_scale, | |
| width=width, | |
| height=height, | |
| generator=generator | |
| ).images[0] | |
| return result, None | |
| except Exception as e: | |
| return None, f"Error generating image: {str(e)}" | |
| # HTMLテンプレート | |
| index_html = """ | |
| <!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>Image Generator</title> | |
| </head> | |
| <body> | |
| <h1>Generate an Image</h1> | |
| <form action="/generate" method="get"> | |
| <label for="prompt">Prompt:</label> | |
| <input type="text" id="prompt" name="prompt" required> | |
| <button type="submit">Generate</button> | |
| </form> | |
| </body> | |
| </html> | |
| """ | |
| # ルートページ | |
| def index(): | |
| return render_template_string(index_html) | |
| # 画像生成エンドポイント | |
| def generate_image(): | |
| 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", 512)) | |
| height = int(request.args.get("height", 512)) | |
| image, error = generate_image_with_diffusers(prompt, negative_prompt, steps, cfg_scale, seed, 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) | |