Update app.py

app.py

@@ -1,13 +1,21 @@
 import json
+
+import torch
+
 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 PIL import Image
 
 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
+from huggan.pytorch.cyclegan.modeling_cyclegan import GeneratorResNet
+from torchvision import transforms as T
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize, RandomCrop, RandomHorizontalFlip
+from torchvision.utils import make_grid
 
 hfapi = HfApi()
 model_names = [model.modelId[model.modelId.index("/") + 1:] for model in hfapi.list_models(author="huggingnft")]
@@ -33,7 +41,7 @@ INTERPOLATION_TEXT = "Text about Interpolation"
 COLLECTION2COLLECTION_TEXT = "Text about Collection2Collection"
 
 STOPWORDS = ["-old"]
-COLLECTION2COLLECTION_KEYS = ["_2_"]
+COLLECTION2COLLECTION_KEYS = ["__2__"]
 
 
 def load_lightweight_model(model_name):
@@ -61,6 +69,11 @@ def clean_models(model_names, stopwords):
             cleaned_model_names.append(model_name)
     return cleaned_model_names
 
+def get_concat_h(im1, im2):
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
 
 model_names = clean_models(model_names, STOPWORDS)
 
@@ -141,7 +154,7 @@ if choose == "Generate image":
                     nrow=nrows,
                     checkpoint=-1,
                     types=generation_type
-                )
+                )[0]
             )
 
 if choose == "Interpolation":
@@ -184,13 +197,75 @@ if choose == "Interpolation":
 
 if choose == "Collection2Collection":
     st.title(choose)
-    st.markdown(INTERPOLATION_TEXT)
+    st.markdown(COLLECTION2COLLECTION_TEXT)
 
     model_name = st.selectbox(
         'Choose model:',
         set(model_names) - set(clean_models(model_names, COLLECTION2COLLECTION_KEYS))
     )
+    nrows = st.number_input("Number of images to generate:",
+                            min_value=1,
+                            max_value=10,
+                            step=1,
+                            value=1,
+                            )
     generate_image_button = st.button("Generate")
 
     if generate_image_button:
-        st.markdown("generating Collection2Collection")
+        n_channels = 3
+
+        image_size = 256
+
+        input_shape = (image_size, image_size)
+
+        transform = Compose([
+            T.ToPILImage(),
+            T.Resize(input_shape),
+            ToTensor(),
+            Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ])
+
+        # generator = modeling_dcgan.Generator.from_pretrained("huggingnft/cryptopunks")
+        with st.spinner(text=f"Downloading selected model..."):
+            translator = GeneratorResNet.from_pretrained(f'huggingnft/{model_name}',
+                                                         input_shape=(n_channels, image_size, image_size),
+                                                         num_residual_blocks=9)
+
+        z = torch.randn(nrows, 100, 1, 1)
+
+        with st.spinner(text=f"Downloading selected model..."):
+            model = load_lightweight_model(f"huggingnft/{model_name.split('__2__')[0]}")
+
+        with st.spinner(text=f"Generating input images..."):
+            punks = model.generate_app(
+                num=timestamped_filename(),
+                nrow=4,
+                checkpoint=-1,
+                types="default"
+            )[1]
+
+        pipe_transform = T.Resize((256, 256))
+
+        input = pipe_transform(punks)
+
+        with st.spinner(text=f"Generating output images..."):
+            output = translator(input)
+
+        out_img = make_grid(output,
+                            nrow=4, normalize=True)
+
+        # out_img = make_grid(punks,
+        # nrow=8, normalize=True)
+
+        out_transform = Compose([
+            T.ToPILImage()
+        ])
+
+        results = []
+
+        for out_punk, out_ape in zip(input, output):
+            results.append(
+                get_concat_h(out_transform(make_grid(out_punk, nrow=1, normalize=True)), out_transform(make_grid(out_ape, nrow=1, normalize=True)))
+            )
+        for result in results:
+            st.image(result)