Update app.py

app.py

@@ -1,121 +1,75 @@
-import subprocess
-from pathlib import Path
-import einops
-import numpy as np
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-from torch import nn
-from torchvision.utils import save_image
+import json
+from huggingnft.lightweight_gan.train import timestamped_filename
+from streamlit_option_menu import option_menu
+
+from huggingface_hub import hf_hub_download, file_download
+
 from huggingface_hub.hf_api import HfApi
 import streamlit as st
+from huggingnft.lightweight_gan.lightweight_gan import Generator, LightweightGAN, evaluate_in_chunks, Trainer
+from accelerate import Accelerator
 
 hfapi = HfApi()
+model_names = [model.modelId[model.modelId.index("/") + 1:] for model in hfapi.list_models(author="huggingnft")]
 
+# streamlit-option-menu
+# st.set_page_config(page_title="Sharone's Streamlit App Gallery", page_icon="", layout="wide")
 
-class Generator(nn.Module):
-    def __init__(self, num_channels=4, latent_dim=100, hidden_size=64):
-        super(Generator, self).__init__()
-        self.model = nn.Sequential(
-            # input is Z, going into a convolution
-            nn.ConvTranspose2d(latent_dim, hidden_size * 8, 4, 1, 0, bias=False),
-            nn.BatchNorm2d(hidden_size * 8),
-            nn.ReLU(True),
-            # state size. (hidden_size*8) x 4 x 4
-            nn.ConvTranspose2d(hidden_size * 8, hidden_size * 4, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(hidden_size * 4),
-            nn.ReLU(True),
-            # state size. (hidden_size*4) x 8 x 8
-            nn.ConvTranspose2d(hidden_size * 4, hidden_size * 2, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(hidden_size * 2),
-            nn.ReLU(True),
-            # state size. (hidden_size*2) x 16 x 16
-            nn.ConvTranspose2d(hidden_size * 2, hidden_size, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(hidden_size),
-            nn.ReLU(True),
-            # state size. (hidden_size) x 32 x 32
-            nn.ConvTranspose2d(hidden_size, num_channels, 4, 2, 1, bias=False),
-            nn.Tanh()
-            # state size. (num_channels) x 64 x 64
-        )
+# sysmenu = '''
+# <style>
+# #MainMenu {visibility:hidden;}
+# footer {visibility:hidden;}
+# '''
+# st.markdown(sysmenu,unsafe_allow_html=True)
 
-    def forward(self, noise):
-        pixel_values = self.model(noise)
-
-        return pixel_values
-
-
-@torch.no_grad()
-def interpolate(model, save_dir='./lerp/', frames=100, rows=8, cols=8):
-    save_dir = Path(save_dir)
-    save_dir.mkdir(exist_ok=True, parents=True)
-
-    z1 = torch.randn(rows * cols, 100, 1, 1)
-    z2 = torch.randn(rows * cols, 100, 1, 1)
-
-    zs = []
-    for i in range(frames):
-        alpha = i / frames
-        z = (1 - alpha) * z1 + alpha * z2
-        zs.append(z)
-
-    zs += zs[::-1]  # also go in reverse order to complete loop
-
-    frames = []
-    for i, z in enumerate(zs):
-        imgs = model(z)
-
-        save_image(imgs, save_dir / f"{i:03}.png", normalize=True)
-        img = Image.open(save_dir / f"{i:03}.png").convert('RGBA')
-        img.putalpha(255)
-        frames.append(img)
-        img.save(save_dir / f"{i:03}.png")
-    frames[0].save("out.gif", format="GIF", append_images=frames,
-                   save_all=True, duration=100, loop=1)
-
-
-def predict(model_name, choice, seed):
-    try:
-        model = Generator(3)
-        weights_path = hf_hub_download(f'huggingnft/{model_name}', 'pytorch_model.bin')
-        model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
-    except:
-        model = Generator(4)
-        weights_path = hf_hub_download(f'huggingnft/{model_name}', 'pytorch_model.bin')
-        model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
-    torch.manual_seed(seed)
-
-    if choice == 'interpolation':
-        interpolate(model)
-        return 'out.gif'
-    else:
-        z = torch.randn(64, 100, 1, 1)
-        punks = model(z)
-        save_image(punks, "image.png", normalize=True)
-        img = Image.open(f"image.png").convert('RGBA')
-        img.putalpha(255)
-        img.save("image.png")
-        return 'image.png'
+# # Add a logo (optional) in the sidebar
+# logo = Image.open(r'C:\Users\13525\Desktop\Insights_Bees_logo.png')
+# profile = Image.open(r'C:\Users\13525\Desktop\medium_profile.png')
 
+ABOUT_TEXT = "🤗 Hugging NFT - Generate NFT by OpenSea collection name."
+CONTACT_TEXT = "Here is some contact info"
+GENERATE_IMAGE_TEXT = "Text about generation"
+INTERPOLATION_TEXT = "Text about Interpolation"
+COLLECTION2COLLECTION_TEXT = "Text about Collection2Collection"
 
-model_names = [model.modelId[model.modelId.index("/") + 1:] for model in hfapi.list_models(author="huggingnft")]
+STOPWORDS = ["-old"]
+COLLECTION2COLLECTION_KEYS = ["2"]
 
-st.set_page_config(page_title="Hugging NFT")
 
-st.title("Hugging NFT")
-st.sidebar.markdown(
-    """
-<style>
-.aligncenter {
-    text-align: center;
-}
-</style>
-<p class="aligncenter">
-    <img src="https://raw.githubusercontent.com/AlekseyKorshuk/optimum-transformers/master/data/social_preview.png" width="300" />
-</p>
-""",
-    unsafe_allow_html=True,
-)
+def load_lightweight_model(model_name):
+    file_path = file_download.hf_hub_download(
+        repo_id=model_name,
+        filename="config.json"
+    )
+    config = json.loads(open(file_path).read())
+    organization_name, name = model_name.split("/")
+    model = Trainer(**config, organization_name=organization_name, name=name)
+    model.load(use_cpu=True)
+    model.accelerator = Accelerator()
+    return model
+
+
+def clean_models(model_names, stopwords):
+    cleaned_model_names = []
+    for model_name in model_names:
+        clear = True
+        for stopword in stopwords:
+            if stopword in model_name:
+                clear = False
+                break
+        if clear:
+            cleaned_model_names.append(model_name)
+    return cleaned_model_names
+
+
+model_names = clean_models(model_names, STOPWORDS)
+
+with st.sidebar:
+    choose = option_menu("Hugging NFT",
+                         ["About", "Generate image", "Interpolation", "Collection2Collection", "Contact"],
+                         icons=['house', 'camera fill', 'bi bi-youtube', 'book', 'person lines fill'],
+                         menu_icon="app-indicator", default_index=0,
+                         )
 st.sidebar.markdown(
     """
 <style>
@@ -123,11 +77,9 @@ st.sidebar.markdown(
     text-align: center;
 }
 </style>
-
 <p style='text-align: center'>
-<a href="https://github.com/AlekseyKorshuk/huggingnft" target="_blank">GitHub</a>
+<a href="https://github.com/AlekseyKorshuk/huggingnft" target="_blank">Project Repository</a>
 </p>
-
 <p class="aligncenter">
     <a href="https://github.com/AlekseyKorshuk/huggingnft" target="_blank"> 
         <img src="https://img.shields.io/github/stars/AlekseyKorshuk/huggingnft?style=social"/>
@@ -142,73 +94,95 @@ st.sidebar.markdown(
     unsafe_allow_html=True,
 )
 
-st.markdown(
-    "🤗 [Hugging NFT](https://github.com/AlekseyKorshuk/huggingnft) - Generate NFT by OpenSea collection name.")
-
-st.markdown(
-    "🚀️ SN-GAN used to train all models.")
-
-st.markdown(
-    "⁉️ Want to train your model? Check [project repository](https://github.com/AlekseyKorshuk/huggingnft) and make this in few clicks!")
-#
-# st.markdown("🚀 Up to 1ms on Bert-based transformers")
-#
-# st.markdown(
-#     "‼️ NOTE: This Space **does not show** the real power of this project because: low recources, not possbile to optimize models. Check [project repository](https://github.com/AlekseyKorshuk/optimum-transformers) with real bechmarks!")
-
-# st.sidebar.header("Settings:")
-model_name = st.selectbox(
-    'Choose model:',
-    model_names)
-
-output_type = st.selectbox(
-    'Output type:',
-    ['image', 'interpolation'])
-
-seed_value = st.slider("Seed:",
-                       min_value=1,
-                       max_value=1000,
-                       step=1,
-                       value=100,
-                       )
-
-model_html = """
-
-<div class="inline-flex flex-col" style="line-height: 1.5;">
-    <div class="flex">
-        <div
-\t\t\tstyle="display:DISPLAY_1; margin-left: auto; margin-right: auto; width: 92px; height:92px; border-radius: 50%; background-size: cover; background-image: url(&#39;USER_PROFILE&#39;)">
-        </div>
-    </div>
-    <div style="text-align: center; margin-top: 3px; font-size: 16px; font-weight: 800">🤖 HuggingArtists Model 🤖</div>
-    <div style="text-align: center; font-size: 16px; font-weight: 800">USER_NAME</div>
-    <a href="https://genius.com/artists/USER_HANDLE">
-    \t<div style="text-align: center; font-size: 14px;">@USER_HANDLE</div>
-    </a>
-</div>
-"""
-
-if st.button("Run"):
-    with st.spinner(text=f"Generating..."):
-        st.image(predict(model_name, output_type, seed_value))
-    st.subheader("Please star project repository, this space and follow my Twitter:")
-    st.markdown(
-        """
-    <style>
-    .aligncenter {
-        text-align: center;
-    }
-    </style>
-    <p class="aligncenter">
-        <a href="https://github.com/AlekseyKorshuk/huggingnft" target="_blank">
-            <img src="https://img.shields.io/github/stars/AlekseyKorshuk/huggingnft?style=social"/>
-        </a>
-    </p>
-    <p class="aligncenter">
-        <a href="https://twitter.com/alekseykorshuk" target="_blank">
-            <img src="https://img.shields.io/twitter/follow/alekseykorshuk?style=social"/>
-        </a>
-    </p>
-        """,
-        unsafe_allow_html=True,
+if choose == "About":
+    st.title(choose)
+    st.markdown(ABOUT_TEXT)
+
+if choose == "Contact":
+    st.title(choose)
+    st.markdown(CONTACT_TEXT)
+
+if choose == "Generate image":
+    st.title(choose)
+    st.markdown(GENERATE_IMAGE_TEXT)
+
+    model_name = st.selectbox(
+        'Choose model:',
+        clean_models(model_names, COLLECTION2COLLECTION_KEYS)
+    )
+    generation_type = st.selectbox(
+        'Select generation type:',
+        ["default", "ema"]
     )
+
+    nrows = st.number_input("Number of rows:",
+                            min_value=1,
+                            max_value=10,
+                            step=1,
+                            value=8,
+                            )
+    generate_image_button = st.button("Generate")
+
+    if generate_image_button:
+        with st.spinner(text=f"Downloading selected model..."):
+            model = load_lightweight_model(f"huggingnft/{model_name}")
+        with st.spinner(text=f"Generating..."):
+            st.image(
+                model.generate_app(
+                    num=timestamped_filename(),
+                    nrow=nrows,
+                    checkpoint=-1,
+                    types=generation_type
+                )
+            )
+
+if choose == "Interpolation":
+    st.title(choose)
+    st.markdown(INTERPOLATION_TEXT)
+
+    model_name = st.selectbox(
+        'Choose model:',
+        clean_models(model_names, COLLECTION2COLLECTION_KEYS)
+    )
+    nrows = st.number_input("Number of rows:",
+                            min_value=1,
+                            max_value=10,
+                            step=1,
+                            value=1,
+                            )
+
+    num_steps = st.number_input("Number of steps:",
+                                min_value=1,
+                                max_value=1000,
+                                step=1,
+                                value=100,
+                                )
+    generate_image_button = st.button("Generate")
+
+    if generate_image_button:
+        with st.spinner(text=f"Downloading selected model..."):
+            model = load_lightweight_model(f"huggingnft/{model_name}")
+        my_bar = st.progress(0)
+        result = model.generate_interpolation(
+            num=timestamped_filename(),
+            num_image_tiles=nrows,
+            num_steps=num_steps,
+            save_frames=False,
+            progress_bar=my_bar
+        )
+        my_bar.empty()
+        with st.spinner(text=f"Uploading result..."):
+            st.image(result)
+
+if choose == "Collection2Collection":
+    st.title(choose)
+    st.markdown(INTERPOLATION_TEXT)
+
+    model_name = st.selectbox(
+        'Choose model:',
+        set(model_names) - set(clean_models(model_names, COLLECTION2COLLECTION_KEYS))
+    )
+    generate_image_button = st.button("Generate")
+
+    if generate_image_button:
+        st.markdown("generating Collection2Collection")