Update app.py

app.py

@@ -4,10 +4,8 @@ import random, re
 
 app = Flask(__name__)
 
-# モデルキャッシュ用の辞書
 model_cache = {}
 
-# テキスト候補読み込み
 with open("ideas.txt", "r") as f:
     lines = f.readlines()
 
@@ -16,25 +14,25 @@ def get_pipeline(model_name):
         model_cache[model_name] = pipeline('text-generation', model=model_name, tokenizer='gpt2')
     return model_cache[model_name]
 
-def generate_prompts(starting_text, model_name, num_prompts=1):
+def generate_prompts(starting_text, model_name, num_prompts=1, generation_args=None):
     response_list = []
     gpt2_pipe = get_pipeline(model_name)
 
     for _ in range(num_prompts):
-        for count in range(4):  # 最大4回試行
+        for count in range(4):
             seed = random.randint(100, 1000000)
             set_seed(seed)
 
-            # 入力テキストが空ならランダムに選ぶ
             if starting_text == "":
                 starting_text = lines[random.randrange(0, len(lines))].strip().lower().capitalize()
                 starting_text = re.sub(r"[,:\-–.!;?_]", '', starting_text)
 
-            # テキスト生成
-            response = gpt2_pipe(starting_text, max_length=random.randint(60, 90), num_return_sequences=1)
+            response = gpt2_pipe(
+                starting_text,
+                **generation_args  # 各種パラメーターを一括で渡す
+            )
             generated_text = response[0]['generated_text'].strip()
 
-            # テキストをチェック・クリーン
             if generated_text != starting_text and len(generated_text) > (len(starting_text) + 4):
                 cleaned_text = re.sub(r'[^ ]+\.[^ ]+', '', generated_text)
                 cleaned_text = cleaned_text.replace("<", "").replace(">", "")
@@ -43,7 +41,6 @@ def generate_prompts(starting_text, model_name, num_prompts=1):
 
     return response_list[:num_prompts]
 
-# APIエンドポイント
 @app.route('/', methods=['GET'])
 def generate_api():
     starting_text = request.args.get('text', default="", type=str)
@@ -51,13 +48,21 @@ def generate_api():
     model_param = request.args.get('model', default="sd", type=str).lower()
 
     # モデル選択
-    if model_param == "dall":
-        model_name = "Gustavosta/MagicPrompt-Dalle"
-    else:
-        model_name = "Gustavosta/MagicPrompt-Stable-Diffusion"
+    model_name = "Gustavosta/MagicPrompt-Dalle" if model_param == "dall" else "Gustavosta/MagicPrompt-Stable-Diffusion"
 
-    # プロンプト生成
-    results = generate_prompts(starting_text, model_name, num_prompts=num_prompts)
+    # URL パラメータから生成設定を取得
+    generation_args = {
+        "max_length": request.args.get('max_length', default=random.randint(60, 90), type=int),
+        "min_length": request.args.get('min_length', default=0, type=int),
+        "temperature": request.args.get('temperature', default=1.0, type=float),
+        "top_k": request.args.get('top_k', default=50, type=int),
+        "top_p": request.args.get('top_p', default=0.95, type=float),
+        "repetition_penalty": request.args.get('repetition_penalty', default=1.0, type=float),
+        "do_sample": request.args.get('do_sample', default=True, type=lambda v: v.lower() in ['true', '1', 'yes']),
+        "num_return_sequences": 1
+    }
+
+    results = generate_prompts(starting_text, model_name, num_prompts=num_prompts, generation_args=generation_args)
     return jsonify(results)
 
 if __name__ == '__main__':