Update app.py

app.py

@@ -19,7 +19,6 @@ def image_to_base64(image_path):
 st.markdown("""
     <style>
     @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@400;700&display=swap');
-
     /* Apply the font to everything */
     html, body, [class*="st"] {
         font-family: 'Roboto', sans-serif;
@@ -130,7 +129,15 @@ if 'generate' not in st.session_state:
 
 # Inizializza inference_tester solo una volta
 if 'inference_tester' not in st.session_state:
-    st.session_state['inference_tester'] = 1
+    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
 
 # Usa inference_tester dalla sessione
 inference_tester = st.session_state['inference_tester']
@@ -202,12 +209,18 @@ 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 col2:
+    with col3:
         if st.button("Return to the beginning"):
             # Ripristina lo stato della sessione
             st.session_state['step'] = 1
@@ -365,8 +378,79 @@ if st.session_state['step'] == 3:
             st.rerun()
 
 if st.session_state['step'] == 4:
-    st.write("Generation completed successfully!")
-    st.session_state['generate'] = False
+    # 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
 
     if st.button("Return to the beginning"):
         # Ripristina lo stato della sessione