Update app.py

app.py

@@ -10,25 +10,6 @@ import torch
 import torch.nn as nn
 from generatorModel import Generator
 
-noise_dim = 16
-embed_dim = 384
-embed_out_dim = 256
-device = 'cpu'
-
-generator = Generator(channels=3, embed_dim=embed_dim, noise_dim=noise_dim, embed_out_dim=embed_out_dim).to(device)
-
-
-# Path to your .pth file
-gen_weight = 'generator_20240421_3.pth'
-
-# Load the weights
-weights_gen = torch.load(gen_weight, map_location=torch.device(device))
-
-# Apply the weights to your model
-generator.load_state_dict(weights_gen)
-
-model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')
-
 def load_embedding(model):
     # Load your model and other components here
     with open(os.path.join("descriptions.json"), 'r') as file:
@@ -39,7 +20,6 @@ def load_embedding(model):
     return embeddings_list, classes
 
 def generate_image(caption):
-    embeddings, classes = load_embedding(model)
     noise_dim = 16
     results = [(util.pytorch_cos_sim(model.encode(caption, convert_to_tensor=True), embeddings[cls]).item(), cls) for cls in classes]
     sorted_results = sorted(results, key=lambda x: x[0], reverse=True)[:5]
@@ -75,10 +55,27 @@ def generate_image(caption):
 
     return img
 
+noise_dim = 16
+embed_dim = 384
+embed_out_dim = 256
+device = 'cpu'
+
+generator = Generator(channels=3, embed_dim=embed_dim, noise_dim=noise_dim, embed_out_dim=embed_out_dim).to(device)
+
+# Path to .pth file and load the weights
+gen_weight = 'generator_20240421_3.pth'
+weights_gen = torch.load(gen_weight, map_location=torch.device(device))
+generator.load_state_dict(weights_gen)
+
+model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')
+
+embeddings, classes = load_embedding(model)
+
+
 iface = gr.Interface(fn=generate_image,
                      inputs=gr.Textbox(lines=2, placeholder="Enter Caption Here..."),
                      outputs=gr.Image(type="numpy"),
                      title="CUHK Shenzhen Building Text-to-Image Generation",
                      description="Enter a caption of some specific building in CUHK-Shenzhen to generate an image..")
 
-iface.launch(share=True, debug=True)
+iface.launch(debug=True)