Update app.py

app.py

@@ -14,6 +14,7 @@ import numpy as np
 from small_256_model import UNet as small_UNet
 from big_1024_model import UNet as big_UNet
 from CLIP import load as load_clip
+from rich import print as rp
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -44,10 +45,10 @@ def load_model():
         model.to(device)
         model.eval()
         global_model = model
-        print("Model loaded successfully!")
+        rp("Model loaded successfully!")
         return model
     except Exception as e:
-        print(f"Error loading model: {e}")
+        rp(f"Error loading model: {e}")
         model = big_UNet().to(device) if big else small_UNet().to(device)
         global_model = model
         return model
@@ -253,10 +254,12 @@ checkpoint = torch.load(name)
 model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
 model.load_state_dict(checkpoint['model_state_dict'])
 model.eval()
+```
 
-Model Architecture
+## Model Architecture
 
 {str(self.model)} """
+            rp(model_card)
             # Save and upload README
             with open("README.md", "w") as f:
                 f.write(model_card)
@@ -352,8 +355,8 @@ def train_model(epochs):
         
             # Compute image reconstruction loss
             img_loss = criterion(output, original)
-            print(f"Original Prompt Tokens Type: {original_prompt_tokens.dtype}, Shape: {original_prompt_tokens.shape}")
-            print(f"Enhanced Prompt Tokens Type: {enhanced_prompt_tokens.dtype}, Shape: {enhanced_prompt_tokens.shape}")
+            rp(f"Image {i} Loss:{imag_loss}")
+            #print(f"Enhanced Prompt Tokens Type: {enhanced_prompt_tokens.dtype}, Shape: {enhanced_prompt_tokens.shape}")
             # Compute prompt guidance loss (L2 norm between original and enhanced prompt embeddings)
             #prompt_loss = torch.norm(original_prompt_tokens - enhanced_prompt_tokens, p=2)
             
@@ -366,7 +369,7 @@ def train_model(epochs):
 
             if i % 10 == 0:
                 status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {total_loss.item():.8f}"
-                print(status)
+                rp(status)
                 output_text.append(status)
         
         # Push model to Hugging Face Hub at the end of each epoch
@@ -422,7 +425,7 @@ def train_model_old(epochs):
 
             if i % 10 == 0:
                 status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {total_loss.item():.8f}"
-                print(status)
+                rp(status)
                 output_text.append(status)
         
         # Push model to Hugging Face Hub at the end of each epoch