anime-kawai-diffusion

Running

App Files Files Community

soiz commited on Oct 24, 2024

Commit

6a3ef3a

•

1 Parent(s): 95d6160

Update app.py

Browse files

Files changed (1) hide show

app.py +80 -83

app.py CHANGED Viewed

@@ -1,80 +1,53 @@
 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>
 """
@@ -82,27 +55,51 @@ kawai diffusion
 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", "")
-    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)

 from flask import Flask, request, jsonify, send_file, render_template_string, make_response
 from deep_translator import GoogleTranslator
+from PIL import Image
+import torch
+from diffusers import StableDiffusionPipeline
+import random
+import io
+import os
 app = Flask(__name__)
+MODEL_NAME = "Ojimi/anime-kawai-diffusion"
+MODEL_DIR = "./models/anime-kawai-diffusion"  # Directory to store the model
+# Download and load the model at startup
+def load_model():
+    if not os.path.exists(MODEL_DIR):
+        print(f"Downloading the model {MODEL_NAME}...")
+        pipeline = StableDiffusionPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
+        pipeline.save_pretrained(MODEL_DIR)
+    else:
+        print(f"Loading model from {MODEL_DIR}...")
+        pipeline = StableDiffusionPipeline.from_pretrained(MODEL_DIR, torch_dtype=torch.float16)
+    if torch.cuda.is_available():
+        pipeline.to("cuda")
+        print("Model loaded on GPU")
+    else:
+        print("GPU not available. Running on CPU.")
+    return pipeline
+# Load the model once during startup
+pipeline = load_model()
 # HTML template for the index page
 index_html = """
 <!DOCTYPE html>
 <html lang="ja">
+<head>
+    <title>Kawaii Diffusion</title>
+</head>
+<body>
+    <h1>Kawaii Diffusion Image Generator</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 Image</button>
+    </form>
+</body>
 </html>
 """
 def index():
     return render_template_string(index_html)
+# Function to generate image locally
+def generate_image_locally(prompt, steps=35, cfg_scale=7, width=512, height=512, seed=-1):
+    # Translate prompt from Russian to English
+    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
+    print(f'Translated prompt: {prompt}')
+    # Set a random seed if not provided
+    generator = torch.manual_seed(seed if seed != -1 else random.randint(1, 1_000_000))
+    # Generate the image using the loaded pipeline
+    image = pipeline(prompt, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator).images[0]
+    return image
 @app.route('/generate', methods=['GET'])
 def generate_image():
+    try:
+        prompt = request.args.get("prompt", "")
+        steps = int(request.args.get("steps", 35))
+        cfg_scale = float(request.args.get("cfgs", 7))
+        width = int(request.args.get("width", 512))
+        height = int(request.args.get("height", 512))
+        seed = int(request.args.get("seed", -1))
+        # Generate the image locally
+        image = generate_image_locally(prompt, steps, cfg_scale, width, height, seed)
+        # Save the image to a BytesIO object
+        img_bytes = io.BytesIO()
+        image.save(img_bytes, format='PNG')
+        img_bytes.seek(0)
+        return send_file(img_bytes, mimetype='image/png')
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+# Content-Security-Policyヘッダーを設定する
+@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
 if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)