streamlit app を作成

.gitignore

@@ -2,5 +2,6 @@
 *.png
 .ipynb_checkpoints/
 __pycache__/
+.env
 
 !/sample_output_images/**

app.py

@@ -0,0 +1,152 @@
+import io
+import os
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import streamlit as st
+import torch
+from diffusers import StableDiffusionPipeline
+from dotenv import load_dotenv
+from huggingface_hub import notebook_login
+from PIL import Image
+
+load_dotenv(".env")
+
+ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN")
+
+sys.path.append("./")
+from simulation import *
+
+# シード値の固定
+SEED = 42
+np.random.seed(seed=SEED)
+random.seed(SEED)
+
+
+def main():
+
+    # 生息地を表すワード
+    HABITAT_WORDS = " Alien from Mars"
+
+    # パラメーター
+    GENOMS_SIZE = 4  # 遺伝配列 0, 1 のどちらかを要素とした配列のサイズ
+    TOUNAMENT_NUM = 10  # トーナメント方式で競わせる数
+    CROSSOVER_PB = 0.8  # cross over(交差) する確率
+    MUTATION_PB = 0.5  # mutation(突然変異)する確率
+
+    # グローバル変数
+    global best
+
+    POPURATIONS = st.slider(
+        label="人口数",
+        min_value=3,
+        max_value=3000,
+        value=500,
+    )
+
+    NUM_GENERATION = st.slider(
+        label="世代数",
+        min_value=10,
+        max_value=10000,
+        value=1000,
+    )
+
+    # キーワード候補
+    word_dict = {
+        "body_size": ["Fingertip sized", "Palm sized", "", "Tall", "Giant"],
+        "body_hair": ["Bald", "Smooth", "", "Furry", "Very Furry"],
+        "herd_num": ["Lone", "Pair", "", "Herd of", "Swarm of"],
+        "eating": ["No teeth", "Herbivorous", "Omnivorous", "Carnivorous", "Fang"],
+        "body_color": [
+            "Lightest skin",
+            "Lighter skin",
+            "",
+            "Darker skin",
+            "Darkest skin",
+        ],
+        "ferocity": ["Peaceful", "Gentle", "", "Ferocious", "Tyrannical"],
+    }
+
+    if st.button("実行", key="ga"):
+
+        st.write("遺伝アルゴリズムの実行")
+
+        progress_bar_ga = st.progress(0)
+
+        # create first genetarion
+        generation = create_generation(POPURATIONS, GENOMS_SIZE)
+
+        progress_bar_ga.progress(50)
+
+        # アルゴリズムの実行
+        best, worst = ga_solve(
+            generation,
+            NUM_GENERATION,
+            POPURATIONS,
+            TOUNAMENT_NUM,
+            CROSSOVER_PB,
+            MUTATION_PB,
+        )
+
+        progress_bar_ga.progress(100)
+
+        st.write("遺伝アルゴリズム処理の終了")
+
+        st.write("画像生成の実行")
+
+        progress_bar_image = st.progress(0)
+
+        progress_bar_image.progress(0)
+
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "CompVis/stable-diffusion-v1-4", use_auth_token=ACCESS_TOKEN
+        )
+        pipe.enable_attention_slicing()
+
+        progress_bar_image.progress(7)
+
+        device = "gpu" if torch.cuda.is_available() else "cpu"
+
+        print("used device is", device)
+        pipe.to(device)
+
+        # NSFWフィルターの回避
+        def null_safety(images, **kwargs):
+            return images, False
+
+        pipe.safety_checker = null_safety
+
+        last_generation = NUM_GENERATION - 1
+
+        plt.figure(figsize=(8, 8))
+        plt.rcParams["font.size"] = 9
+
+        words = (
+            get_word_for_image_generate(word_dict, best, last_generation)
+            + HABITAT_WORDS
+        )
+
+        image = pipe(words)["sample"][0]
+
+        plt.title(f"{last_generation + 1}th\n{words}.")
+        plt.xticks([])
+        plt.yticks([])
+        plt.imshow(image)
+
+        progress_bar_image.progress(100)
+
+        plt.tight_layout()
+        buf = io.BytesIO()
+        plt.savefig(buf, format="png")
+
+        buf.seek(0)
+        im = Image.open(buf)
+        numpy_image = np.array(im)
+        st.image(numpy_image)
+
+
+if __name__ == "__main__":
+
+    main()

requirements.txt

@@ -5,4 +5,6 @@ pandas==1.3.5
 ftfy==6.1.1
 spacy==3.4.1
 matplotlib==3.5.3
-notebook==6.4.12
+notebook==6.4.12
+streamlit==1.13.0
+python-dotenv==0.21.0