Update app.py

app.py

@@ -16,11 +16,11 @@ class dotdict(dict):
     __delattr__ = dict.__delitem__
 
 class InferRunner:
-    def __init__(self):
+    def __init__(self, device):
         vae_config = json.load(open("ckpts/ldm/vae_config.json"))
         self.vae = AutoencoderKL(**vae_config).to(device)
-        vae_weights = torch.load("ckpts/ldm/pytorch_model_vae.bin", map_location=device)
-        self.vae.load_state_dict(vae_weights)
+        # vae_weights = torch.load("ckpts/ldm/pytorch_model_vae.bin", map_location=device)
+        # self.vae.load_state_dict(vae_weights)
 
         train_args = dotdict(json.loads(open("ckpts/pico_model/summary.jsonl").readlines()[0]))
         self.pico_model = PicoDiffusion(
@@ -39,13 +39,9 @@ def infer(caption, runner):
         wave = runner.vae.decode_to_waveform(mel)[0][:audio_len]
     sf.write(f"synthesized/{caption}.wav", wave, samplerate=16000, subtype='PCM_16')
 
-infer_runner = InferRunner()
-if torch.cuda.is_available():
-    device = "cuda"
-    device_selection = "cuda:0"
-else:
-    device = "cpu"
-    device_selection = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+infer_runner = InferRunner(device)
+
 
 with gr.Blocks() as demo:
     with gr.Row():