add models

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from huggingface_hub import hf_hub_download
 from prediction import run_image_prediction
 import torch
 import torchvision.transforms as T
@@ -9,10 +8,9 @@ from matplotlib import pyplot as plt
 
 
 def gradio_demo(model_name, sequence_input, nucleus_image, protein_image):
-    model = hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="model.ckpt")
-    config = hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="config.yaml")
-    hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="nucleus_vqgan.yaml")
-    hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="threshold_vqgan.yaml")
+    model = f"models/{model_name}.ckpt"
+    config = f"models/{model_name}.yaml"
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     if "Finetuned" in model_name:

models/HPA_480.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14b445eedda4f46e67c7d848df3301363e9034d752591a1275b47eaf483e0cac
+size 4731482483

models/HPA_480.yaml

@@ -0,0 +1,25 @@
+model:
+  learning_rate: 0.0003
+  target: celle_main.CELLE_trainer
+  params:
+    ckpt_path: model.ckpt
+    condition_model_path: 
+    condition_config_path: nucleus_vqgan.yaml
+    vqgan_model_path: 
+    vqgan_config_path: threshold_vqgan.yaml
+    image_key: threshold
+    num_images: 2
+    dim: 480
+    num_text_tokens: 33
+    text_seq_len: 1000
+    depth: 68
+    heads: 16
+    dim_head: 64
+    attn_dropout: 0.1
+    ff_dropout: 0.1
+    attn_types: full
+    rotary_emb: true
+    fixed_embedding: true
+    text_embedding: esm2
+    loss_img_weight: 1
+    loss_cond_weight: 1

models/HPA_Finetuned_480.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:152dfd30c8adfe83868ac0ffea4d075d25d11f24369f34482f79a3db4524d1be
+size 4731482675

models/HPA_Finetuned_480.yaml

@@ -0,0 +1,25 @@
+model:
+  learning_rate: 0.0003
+  target: celle_main.CELLE_trainer
+  params:
+    ckpt_path: model.ckpt
+    condition_model_path: 
+    condition_config_path: nucleus_vqgan.yaml
+    vqgan_model_path: 
+    vqgan_config_path: threshold_vqgan.yaml
+    image_key: threshold
+    num_images: 2
+    dim: 480
+    num_text_tokens: 33
+    text_seq_len: 1000
+    depth: 68
+    heads: 16
+    dim_head: 64
+    attn_dropout: 0.1
+    ff_dropout: 0.1
+    attn_types: full
+    rotary_emb: true
+    fixed_embedding: true
+    text_embedding: esm2
+    loss_img_weight: 1
+    loss_cond_weight: 1

models/nucleus_vqgan.yaml

@@ -0,0 +1,33 @@
+model:
+  base_learning_rate: 4.5e-06
+  target: taming.models.vqgan.VQModel
+  params:
+    image_key: nucleus
+    ckpt_path: 
+    embed_dim: 256
+    n_embed: 512
+    ddconfig:
+      double_z: false
+      z_channels: 256
+      resolution: 256
+      in_channels: 1
+      out_ch: 1
+      ch: 128
+      ch_mult:
+      - 1
+      - 1
+      - 2
+      - 2
+      - 4
+      num_res_blocks: 2
+      attn_resolutions:
+      - 16
+      dropout: 0.0
+    lossconfig:
+      target: taming.modules.losses.vqperceptual.VQLPIPSWithDiscriminator
+      params:
+        disc_conditional: false
+        disc_in_channels: 1
+        disc_start: 50000
+        disc_weight: 0.2
+        codebook_weight: 1.0

models/threshold_vqgan.yaml

@@ -0,0 +1,33 @@
+model:
+  base_learning_rate: 4.5e-05
+  target: taming.models.vqgan.VQModel
+  params:
+    image_key: threshold
+    ckpt_path: 
+    embed_dim: 256
+    n_embed: 512
+    ddconfig:
+      double_z: false
+      z_channels: 256
+      resolution: 256
+      in_channels: 1
+      out_ch: 1
+      ch: 128
+      ch_mult:
+      - 1
+      - 1
+      - 2
+      - 2
+      - 4
+      num_res_blocks: 2
+      attn_resolutions:
+      - 16
+      dropout: 0.0
+    lossconfig:
+      target: taming.modules.losses.vqperceptual.VQLPIPSWithDiscriminator
+      params:
+        disc_conditional: false
+        disc_in_channels: 1
+        disc_start: 100000
+        disc_weight: 0.2
+        codebook_weight: 1.0

requirements.txt

@@ -1,7 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch
 torchvision
-huggingface_hub
 gradio
 OmegaConf
 axial-positional-embedding