Update app.py

app.py

@@ -129,15 +129,7 @@ if 'generate' not in st.session_state:
 
 # Inizializza inference_tester solo una volta
 if 'inference_tester' not in st.session_state:
-    model_load_paths = ['CoDi_encoders.pth', 'CoDi_text_diffuser.pth', 'CoDi_video_diffuser_8frames.pth']
-    st.session_state['inference_tester'] = dani_model(model='thesis_model',
-                                                      data_dir='/mimer/NOBACKUP/groups/snic2022-5-277/dmolino/checkpoints/',
-                                                      pth=model_load_paths, load_weights=False)
-    inference_tester = st.session_state['inference_tester']
-
-    # Caricamento dei pesi Clip, Optimus, Frontal, Lateral e Text una sola volta
-    if 'weights_loaded' not in st.session_state:
-        st.session_state['weights_loaded'] = True  # Indica che i pesi sono stati caricati
+    st.session_state['inference_tester'] = 1
 
 # Usa inference_tester dalla sessione
 inference_tester = st.session_state['inference_tester']
@@ -209,18 +201,12 @@ if st.session_state['step'] == 2:
 
     # Pulsante per provare un esempio
     with col1:
-        if st.button("Inference"):
-            st.session_state['step'] = 3  # Passa al passo 3
-            st.rerun()
-
-    # Pulsante per provare un esempio
-    with col2:
         if st.button("Try an example"):
             st.session_state['step'] = 5  # Passa al passo 5
             st.rerun()
 
     # Pulsante per tornare all'inizio
-    with col3:
+    with col2:
         if st.button("Return to the beginning"):
             # Ripristina lo stato della sessione
             st.session_state['step'] = 1
@@ -378,79 +364,8 @@ if st.session_state['step'] == 3:
             st.rerun()
 
 if st.session_state['step'] == 4:
-    # Costruzione del prompt
-    if st.session_state['generate'] is True:
-        conditioning = []
-        for inp in st.session_state['inputs']:
-            if inp == 'frontal':
-                cim = inference_tester.net.clip_encode_vision(st.session_state['frontal'], encode_type='encode_vision').to(device)
-                uim = inference_tester.net.clip_encode_vision(torch.zeros_like(st.session_state['frontal']).to(device),
-                                                              encode_type='encode_vision').to(device)
-                conditioning.append(torch.cat([uim, cim]))
-            elif inp == 'lateral':
-                cim = inference_tester.net.clip_encode_vision(st.session_state['lateral'], encode_type='encode_vision').to(device)
-                uim = inference_tester.net.clip_encode_vision(torch.zeros_like(st.session_state['lateral']).to(device),
-                                                              encode_type='encode_vision').to(device)
-                conditioning.append(torch.cat([uim, cim]))
-            elif inp == 'text':
-                ctx = inference_tester.net.clip_encode_text(1 * [st.session_state['report']], encode_type='encode_text').to(device)
-                utx = inference_tester.net.clip_encode_text(1 * [""], encode_type='encode_text').to(device)
-                conditioning.append(torch.cat([utx, ctx]))
-
-        # Costruzione delle shapes
-        shapes = []
-        for out in st.session_state['outputs']:
-            if out == 'frontal' or out == 'lateral':
-                shape = [1, 4, 256 // 8, 256 // 8]
-                shapes.append(shape)
-            elif out == 'text':
-                shape = [1, 768]
-                shapes.append(shape)
-
-        progress_bar = st.progress(0)
-
-        # Inferenza
-        z, _ = inference_tester.sampler.sample(
-            steps=50,
-            shape=shapes,
-            condition=conditioning,
-            unconditional_guidance_scale=7.5,
-            xtype=st.session_state['outputs'],
-            condition_types=st.session_state['inputs'],
-            eta=1,
-            verbose=False,
-            mix_weight={'lateral': 1, 'text': 1, 'frontal': 1},
-            progress_bar=progress_bar)
-
-        # Decoder e visualizzazione dei risultati
-        output_cols = st.columns(len(st.session_state['outputs']))
-
-        # Definire due colonne per le immagini
-        col1, col2 = st.columns(2)
-
-        # Iterare sugli output e assegnare le immagini alle colonne corrispondenti
-        for i, out in enumerate(st.session_state['outputs']):
-            if out == 'frontal':
-                x = inference_tester.net.autokl_decode(z[i])
-                x = torch.clamp((x[0] + 1.0) / 2.0, min=0.0, max=1.0)
-                im = x[0].cpu().numpy()
-                with col1:  # Mostrare la frontal image nella prima colonna
-                    st.image(im, caption="Generated Frontal Image")
-            elif out == 'lateral':
-                x = inference_tester.net.autokl_decode(z[i])
-                x = torch.clamp((x[0] + 1.0) / 2.0, min=0.0, max=1.0)
-                im = x[0].cpu().numpy()
-                with col2:  # Mostrare la lateral image nella seconda colonna
-                    st.image(im, caption="Generated Lateral Image")
-            elif out == 'text':
-                x = inference_tester.net.optimus_decode(z[i], max_length=100)
-                x = [a.tolist() for a in x]
-                rec_text = [inference_tester.net.optimus.tokenizer_decoder.decode(a) for a in x]
-                rec_text = rec_text[0].replace('<BOS>', '').replace('<EOS>', '')
-                st.write(f"Generated Report: {rec_text}")
-
-        st.write("Generation completed successfully!")
-        st.session_state['generate'] = False
+    st.write("Generation completed successfully!")
+    st.session_state['generate'] = False
 
     if st.button("Return to the beginning"):
         # Ripristina lo stato della sessione
@@ -564,4 +479,4 @@ if st.session_state['step'] == 5:
         st.session_state['frontal_file'] = None
         st.session_state['lateral_file'] = None
         st.session_state['report'] = ""
-        st.rerun()
+        st.rerun()