updated the model loader and readme

.github/workflows/hugging-face.yaml

@@ -12,6 +12,27 @@ jobs:
         with:
           fetch-depth: 0
           lfs: true
+
+      - name: Update README.md
+        run: |
+          tmp_file=$(mktemp)
+          echo "---" >> $tmp_file
+          echo "title: Ve Gans" >> $tmp_file
+          echo "emoji: 📈" >> $tmp_file
+          echo "colorFrom: yellow" >> $tmp_file
+          echo "colorTo: gray" >> $tmp_file
+          echo "sdk: streamlit" >> $tmp_file
+          echo "sdk_version: 1.29.0" >> $tmp_file
+          echo "app_file: space.py" >> $tmp_file
+          echo "pinned: false" >> $tmp_file
+          echo "license: openrail" >> $tmp_file
+          echo "---" >> $tmp_file
+          echo "" >> $tmp_file
+          cat README.md >> $tmp_file
+          mv $tmp_file README.md
+          git add README.md
+          git commit -m "Updated README.md"
+
       - name: Push to hub
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}

README.md

@@ -1,15 +1,3 @@
----
-title: Ve Gans
-emoji: 👀
-colorFrom: yellow
-colorTo: green
-sdk: streamlit
-sdk_version: 1.29.0
-app_file: app.py
-pinned: false
-license: openrail
----
-
 # ve-gans: Image Generation with GANs using PyTorch
 
 ## Overview

requirements.txt

@@ -0,0 +1,5 @@
+streamlit
+torch
+torchvision
+Pillow
+numpy

sample.py

@@ -4,6 +4,6 @@ vegans = Vegans
 
 vegans.load_pretrained()
 
-text = 'something'
+text = "something"
 
-vegans.generate(text)
+vegans.generate(text)

setup.py

@@ -1,21 +1,21 @@
 from setuptools import setup, find_packages
 
-with open('requirements.txt') as f:
+with open("requirements.txt") as f:
     requirements = f.read().splitlines()
 
 setup(
-    name='ve-gans',
-    version='0.1',
+    name="ve-gans",
+    version="0.1",
     packages=find_packages(),
     install_requires=requirements,
-    author='Zai',
-    author_email='[email protected]',
-    description='Floorplan generation model with pytorch',
-    long_description='Detailed description of your project',
-    url='https://github.com/zaibutcooler/ve-gans',
+    author="Zai",
+    author_email="[email protected]",
+    description="Floorplan generation model with pytorch",
+    long_description="Detailed description of your project",
+    url="https://github.com/zaibutcooler/ve-gans",
     classifiers=[
-        'Programming Language :: Python :: 3',
-        'License :: OSI Approved :: MIT License',
-        'Operating System :: OS Independent',
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
     ],
 )

space.py

@@ -1,14 +1,70 @@
 import streamlit as st
 import torch
-from torchvision.utils import make_grid
-from torchvision.transforms import ToPILImage
-from .vegans import Vegans
-
-
-def main():
-    st.title("Image Generation")
-    st.write("Made with GANS from scratch")
-    
-    
-if __name__ == '__main__':
-    main()
+import torch.nn as nn
+import torch.optim as optim
+from PIL import Image
+import numpy as np
+
+# GAN model architecture
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        # TO DO: implement generator architecture
+
+    def forward(self, z):
+        # TO DO: implement generator forward pass
+        pass
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        # TO DO: implement discriminator architecture
+
+    def forward(self, x):
+        # TO DO: implement discriminator forward pass
+        pass
+
+# Floorplan generator model architecture
+class FloorplanGenerator(nn.Module):
+    def __init__(self):
+        super(FloorplanGenerator, self).__init__()
+        # TO DO: implement floorplan generator architecture
+
+    def forward(self, img):
+        # TO DO: implement floorplan generator forward pass
+        pass
+
+# Load pre-trained models
+@st.cache
+def load_models():
+    generator = Generator()
+    discriminator = Discriminator()
+    floorplan_generator = FloorplanGenerator()
+    # TO DO: load pre-trained model weights
+    return generator, discriminator, floorplan_generator
+
+# Streamlit app
+st.title("GAN Image Generation and Floorplan App")
+
+# Load models
+generator, discriminator, floorplan_generator = load_models()
+
+# Image generation
+st.header("Image Generation")
+z_dim = 100
+noise = torch.randn(1, z_dim)
+generated_img = generator(noise)
+st.image(generated_img, caption="Generated Image")
+
+# Floorplan generation
+st.header("Floorplan Generation")
+img_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+if img_file:
+    img = Image.open(img_file)
+    img = np.array(img) / 255.0
+    floorplan = floorplan_generator(torch.tensor(img))
+    st.image(floorplan, caption="Generated Floorplan")
+
+# Run the app
+if __name__ == "__main__":
+    st.write("App is running!")

tests/test_dataset.py

@@ -1,7 +1,14 @@
 import unittest
+from vegans.dataset import TrainingSet
+
 
 class TestTraining(unittest.TestCase):
-    pass
+    def test_download_dataset(self):
+        data = TrainingSet()
+        data.download_dataset()
+        assert data.images is not None
+        assert data.labels is not None
+
 
-if __name__ == '__main__':
-    unittest.main()
+if __name__ == "__main__":
+    unittest.main()

tests/test_generating.py

@@ -1,7 +1,17 @@
 import unittest
+from vegans import Vegans
+
 
 class TestTraining(unittest.TestCase):
-    pass
+    def test_loading_pretrained(self):
+        model = Vegans()
+        model.load_pretrained()
+
+    def test_generator(self):
+        model = Vegans()
+        print("TODO, make generator")
+        pass
+
 
-if __name__ == '__main__':
-    unittest.main()
+if __name__ == "__main__":
+    unittest.main()

tests/test_space.py

@@ -0,0 +1,11 @@
+import unittest
+import space
+
+
+class TestTraining(unittest.TestCase):
+    def test_running(self):
+        pass
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_training.py

@@ -1,7 +1,12 @@
 import unittest
+from vegans import Vegans
+
 
 class TestTraining(unittest.TestCase):
-    pass
+    def test_training(self):
+        model = Vegans()
+        model.train(upload=False, real=False)
+
 
-if __name__ == '__main__':
-    unittest.main()
+if __name__ == "__main__":
+    unittest.main()

train.py

@@ -8,4 +8,3 @@ vegans.train()
 
 # or you can simply just
 # vegans.load_pretrained()
-

vegans/config.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel
+
+
+class Config(BaseModel):
+    learning_rate: float = 0.0002  # learning rate for generator and discriminator
+    num_epoch: int = 100  # number of epochs to train
+    betas: tuple[float, float] = (0.5, 0.999)  # beta1 and beta2 for Adam optimizer
+    batch_size: int = 32  # batch size for training
+    image_size: int = 64  # size of input images
+    channels: int = 3  # number of color channels in input images
+    n_critic: int = 5  # number of critic iterations per generator iteration
+    save_interval: int = 100  # interval for saving model checkpoints
+    sample_interval: int = 100  # interval for generating sample images

vegans/dataset.py

@@ -1,26 +1,36 @@
 from torch.utils.data import Dataset
 from torchvision import transforms as tran
 from datasets import load_dataset
+from .utils import log
+from .config import Config
 
 
+class TrainData(Dataset):
+    def __init__(self, config: Config, gray_scale=True):
+        super().__init__()
 
-class TrainingSet(Dataset):
-    def __init__(self,image_size=(128,128),gray_scale=True):
-        super().__init__()        
-        self.transforms = tran.Compose([
-            tran.Resize(image_size),
-            tran.ToTensor(),
-            tran.Normalize([0.5],[0.5]),
-            tran.Grayscale() if gray_scale==True else None,
-        ])
-
+        self.transforms = tran.Compose(
+            [
+                tran.Resize((config.image_size, config.image_size)),
+                tran.ToTensor(),
+                tran.Normalize([0.5], [0.5]),
+                tran.Grayscale() if gray_scale else None,
+            ]
+        )
         self.images = None
         self.labels = None
- 
+        self.download_dataset()
+
+    def download_dataset(self):
+        dataset = load_dataset("zaibutcooler/archi-vault")
+        self.images = dataset["train"]["images"]
+        self.labels = dataset["train"]["labels"]
+        log("Successfully loaded the dataset")
+
     def __getitem__(self, index):
         image = self.transforms(self.images[index])
         label = self.labels[index]
-        return image,label
-    
+        return image, label
+
     def __len__(self):
         return len(self.images)

vegans/discriminator.py

@@ -1,13 +1,14 @@
 import torch.nn as nn
 import torch.nn.functional as F
-import torch
+from .config import Config
+from huggingface_hub import PyTorchModelHubMixin
 
-class Discriminator(nn.Module):
-  def __init__(self):
-    super(Discriminator, self).__init__()
-    self.model = nn.Sequential(
 
-    )
+class Discriminator(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, config: Config):
+        super(Discriminator, self).__init__()
+        self.config = config
+        self.model = nn.Sequential(nn.Linear())
 
-  def forward(self,x,y):
-    return x
+    def forward(self, x, y):
+        return x

vegans/generator.py

@@ -2,15 +2,18 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch
 
-class Generator(nn.Module):
-  def __init__(self):
-    super(Generator, self).__init__()
-    def block(in_feat,out_feat,norm=False):
-      pass
-    
-    self.model = nn.Sequential(
 
-    )
+from huggingface_hub import PyTorchModelHubMixin
 
-  def forward(self,x,y):
-    pass
+
+class Generator(nn.Module, PyTorchModelHubMixin):
+    def __init__(self):
+        super(Generator, self).__init__()
+
+        def block(in_feat, out_feat, norm=False):
+            pass
+
+        self.model = nn.Sequential()
+
+    def forward(self, x, y):
+        pass

vegans/utils.py

@@ -1,10 +1,31 @@
+from PIL import Image
+import os
+
+
 def display_image():
     pass
 
-def save_image(image,location):
-    pass
+
+def save_image(tensor, location="./"):
+
+    image = Image.fromarray(tensor)
+
+    if not os.path.exists(location):
+        os.makedirs(location)
+
+    image.save(os.path.join(location, "image.jpg"))
+
 
 def log(text):
+    print("")
     print("#############################################\n")
     print(f"{text}\n")
     print("#############################################\n")
+
+
+import torch
+
+
+def generate_noise(batch_size, z_dim):
+
+    return torch.randn(batch_size, z_dim)

vegans/vegans.py

@@ -1,54 +1,95 @@
 # initialize vegans
 import torch
-from generator import Generator
-from discriminator import Discriminator
-from discriminator import Discriminator
-from utils import save_image,display_image,log
-from dataset import TrainingSet
+import huggingface_hub
+
+from .generator import Generator
+from .discriminator import Discriminator
+from .utils import save_image, display_image, log, generate_noise
+from .dataset import TrainingSet
+from .config import Config
 
 
 class Vegans:
-    def __init__(self,lr=0.01,b1=0.02,b2=0.02):
-        self.learning_rate = lr
-        self.num_epoch = 0
-        self.betas = (b1,b2)
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    def __init__(self, config: Config, dataset):
+        assert config is not None
+        self.config = config
+        self.dataset = dataset
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.generator = Generator().to(self.device)
         self.discriminator = Discriminator().to(self.device)
-        self.dataset = TrainingSet()
-        self.g_optim = torch.optim.Adam(self.generator.parameters(),lr=self.learning_rate,betas=self.betas)
-        self.d_optim = torch.optim.Adam(self.discriminator.parameters(),lr=self.learning_rate,betas=self.betas)
+        self.g_optim = torch.optim.Adam(
+            self.generator.parameters(),
+            lr=self.config.learning_rate,
+            betas=self.config.betas,
+        )
+        self.d_optim = torch.optim.Adam(
+            self.discriminator.parameters(),
+            lr=self.config.learning_rate,
+            betas=self.config.betas,
+        )
+
+    def train(self, upload=False, real=True):
 
-    def train(self):
         loss_fn = torch.nn.BCELoss()
         log("Started Training Loop")
+        if real:
+            dataset = self.dataset
+        else:
+            dataset = self._dummy_dataset()
 
         for epoch in range(self.num_epoch):
-            for i,(image,label) in enumerate(self.dataset):
-                
-                self.g_optim.zero_grad()
+            for i, (image, label) in enumerate(dataset):
 
-                loss = loss_fn(0,0)
+                # Train the generator
+                self.g_optim.zero_grad()
 
+                loss = loss_fn(0, 0)
 
+                # Train the discriminator
                 self.d_optim.zero_grad()
 
-            print(f"Epoch {epoch} done. Loss is {loss.item()}`")
+            log(f"Epoch {epoch} done. Loss is {loss.item()}")
 
         log("Finish Training")
+        if upload:
+            self.generator.save_pretrained("ve-gans")
+            self.generator.push_to_hub("ve-gans")
 
-    def generate(self,label):
-        # TODO
-        noise = 0
-        output = self.generator(noise,label)
+    def eval(self):
+        pass
 
-    def load_pretrained(self,checkpoint):
+    def _dummy_dataset(self):
+        dummy_images = []
+        dummy_labels = []
+        for i in range(100):
+            image = torch.randn(
+                1, self.config.channels, self.config.image_size, self.config.image_size
+            )
+            label = torch.randint(0, 2, (1,))
+            dummy_images.append(image)
+            dummy_labels.append(label)
+        return [(image, label) for image, label in zip(dummy_images, dummy_labels)]
 
+    def generate(self, label, save=True):
         # TODO
-        saved_checkpoint = checkpoint
-
-        model_checkpoint = torch.load('')
-        self.generator.load_state_dict(model_checkpoint)
+        noise = generate_noise()
+        output = self.generator(noise, label)
+        if save:
+            save_image(output)
+        return output
 
-        log("Successfully loaded model")
+    def save_pretrained(self, name="ve-gans"):
+        print("Uploading model...")
+        self.model.save_pretrained(name)
+        print(f"Model saved locally as '{name}'")
+        self.model.push_to_hub(name)
+        print(f"Model '{name}' uploaded to the Hugging Face Model Hub")
 
+    def load_pretrained(self, model_id="zaibutcooler/ve-gans"):
+        print("Loading model...")
+        model = model.from_pretrained(model_id)
+        print(f"Model '{model_id}' loaded successfully")
+        return model
+    
+    def huggingface_login(self,token):
+        huggingface_hub.login(token)

version.py

@@ -1 +1 @@
-__version__ = "20231117"
+__version__ = "20231117"