Update app.py

app.py

@@ -1,6 +1,5 @@
 from flask import Flask, request, jsonify, send_file, render_template_string, make_response
-import torch
-from diffusers import StableDiffusionPipeline
+import requests
 import io
 import random
 from PIL import Image
@@ -8,49 +7,63 @@ 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"
+API_URL = "https://api-inference.huggingface.co/models/Ojimi/anime-kawai-diffusion"
+timeout = 3000  # タイムアウトを300秒に設定
 
-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):
+# 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"
 
-    # Translate the prompt from Russian to English (if needed)
+    key = random.randint(0, 999)
+    
+    # Translate the prompt from Russian to English if necessary
     prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
-    print(f"Translated prompt: {prompt}")
+    print(f'Generation {key} translation: {prompt}')
 
-    # Add artistic flair to the prompt
+    # Add some extra 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
+    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'; "
@@ -61,17 +74,7 @@ def add_security_headers(response):
 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>
+kawai diffusion
 </html>
 """
 
@@ -80,26 +83,26 @@ def index():
     return render_template_string(index_html)
 
 @app.route('/generate', methods=['GET'])
-def generate():
+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))
+    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))
 
-    # Generate the image locally
-    image, error = generate_image_locally(prompt, negative_prompt, steps, cfg_scale, seed, width, height)
+    image, error = query(prompt, negative_prompt, steps, cfg_scale, sampler, seed, strength, 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)
+    app.run(host='0.0.0.0', port=7860)