Image generation error fix

app.py

@@ -97,14 +97,14 @@ def display_sample_pytorch(seed, truncation, directions, distances, scale, start
     else:
         w_numpy = [x.cpu().detach().numpy() for x in w]
         w = [np.expand_dims(x, 0) for x in w_numpy]
-        w = [torch.from_numpy(x).to(device) for x in w]
         #w = [x.unsqueeze(0) for x in w]
 
     
     for l in range(start, end):
       for i in range(len(directions)):
         w[l] = w[l] + directions[i] * distances[i] * scale
-    
+
+    w = [torch.from_numpy(x).to(device) for x in w]
     torch.cuda.empty_cache()
     #save image and display
     out = model.sample(w)
@@ -255,9 +255,9 @@ if submit_button:  # Execute when the submit button is pressed
     w2 = clip_optimized_latent(text2, seed2, iters)
     st.session_state['w2-np'] = w2
 
-input_im, output_im = generate_image(content, style, truncation, c0, c1, c2, c3, c4, c5, c6, start_layer, end_layer,st.session_state['w1-np'],st.session_state['w2-np'])
-st.image(input_im, caption="Input Image")
-st.image(output_im, caption="Output Image")
-
-
-        
+try:
+    input_im, output_im = generate_image(content, style, truncation, c0, c1, c2, c3, c4, c5, c6, start_layer, end_layer,st.session_state['w1-np'],st.session_state['w2-np'])
+    st.image(input_im, caption="Input Image")
+    st.image(output_im, caption="Output Image")
+except:
+    pass