|
import streamlit as st |
|
import pandas as pd |
|
import networkx as nx |
|
import os |
|
import pickle |
|
import tqdm |
|
from analysis import build_graph, parse_page |
|
|
|
def clean(results): |
|
new = {} |
|
for k in results: |
|
if results[k] and len(results[k]) > 0: |
|
new[k] = results[k] |
|
return new |
|
|
|
|
|
if "B_degree_threshold" not in st.session_state: |
|
st.session_state.B_degree_threshold = 10 |
|
if "B" not in st.session_state: |
|
if not os.path.exists('data.pkl'): |
|
page_folder = 'pages' |
|
pages = os.listdir(page_folder) |
|
results = {} |
|
|
|
for p in tqdm.tqdm(pages): |
|
try: |
|
results[p] = parse_page(os.path.join(page_folder, p)) |
|
except Exception as e: |
|
pass |
|
|
|
with open('data.pkl', 'wb') as f: |
|
pickle.dump(results, f) |
|
|
|
else: |
|
with open('data.pkl', 'rb') as f: |
|
results = pickle.load(f) |
|
|
|
st.session_state.results = clean(results) |
|
st.session_state.B = build_graph(st.session_state.results, st.session_state.B_degree_threshold) |
|
|
|
|
|
|
|
|
|
def main(): |
|
st.title("SD BaseModel Lora Connections") |
|
|
|
|
|
B_degree_threshold = st.sidebar.slider("Select Degree Threshold", 1, 100, 10) |
|
|
|
|
|
if B_degree_threshold != st.session_state.B_degree_threshold: |
|
st.session_state.B_degree_threshold = B_degree_threshold |
|
st.session_state.B = build_graph(st.session_state.results, B_degree_threshold) |
|
|
|
st.sidebar.write(f"There are {len(st.session_state.B)} nodes analyzed.") |
|
|
|
|
|
model_nodes = {n for n, d in st.session_state.B.nodes(data=True) if d['bipartite']==0} |
|
lora_nodes = set(st.session_state.B) - model_nodes |
|
|
|
|
|
sorted_models = sorted(model_nodes, key=lambda x: st.session_state.B.degree(x), reverse=True) |
|
sorted_loras = sorted(lora_nodes, key=lambda x: st.session_state.B.degree(x), reverse=True) |
|
|
|
|
|
selected_model = st.selectbox("Select Model (sorted by degree)", sorted_models) |
|
if selected_model: |
|
loras_for_model = list(st.session_state.B.neighbors(selected_model)) |
|
page_names_for_model = [st.session_state.B[selected_model][lora]['page'] for lora in loras_for_model] |
|
page_names_for_model = ['https://civitai.com/images/'+page for page in page_names_for_model] |
|
|
|
|
|
df = pd.DataFrame({"Lora Names": loras_for_model, "Image Link": page_names_for_model}) |
|
df["Image Link"] = df["Image Link"].apply(lambda x: f'<a href="{x}" target="_blank">{x}</a>') |
|
st.markdown(df.to_html(escape=False, index=False), unsafe_allow_html=True) |
|
|
|
|
|
selected_lora = st.selectbox("Select Lora (sorted by degree)", sorted_loras) |
|
if selected_lora: |
|
models_for_lora = list(st.session_state.B.neighbors(selected_lora)) |
|
page_names_for_lora = [st.session_state.B[model][selected_lora]['page'] for model in models_for_lora] |
|
page_names_for_lora = ['https://civitai.com/images/'+page for page in page_names_for_lora] |
|
|
|
|
|
df = pd.DataFrame({"Model Names": models_for_lora, "Image Link": page_names_for_lora}) |
|
df["Image Link"] = df["Image Link"].apply(lambda x: f'<a href="{x}" target="_blank">{x}</a>') |
|
st.markdown(df.to_html(escape=False, index=False), unsafe_allow_html=True) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|