Pushing deployment to space

generate_report.py

@@ -12,32 +12,58 @@ from huggingface_hub import hf_hub_download
 # from src.models.cnn_encoder import 
 # from src.models.distil_gpt2 import DistilGPT2
 # from src.configs import argHandler
-
 FLAGS = argHandler()
-FLAGS.setDefaults()
-tokenizer_wrapper = TokenizerWrapper( FLAGS.csv_label_columns[0], FLAGS.max_sequence_length, FLAGS.tokenizer_vocab_size)
+def init_model():
+    global tokenizer_wrapper, encoder, decoder, optimizer
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    print("✅ Initializing model components...")
+
+    from configs import argHandler
+    
+    FLAGS.setDefaults()
+
+    tokenizer_wrapper = TokenizerWrapper(
+        FLAGS.csv_label_columns[0],
+        FLAGS.max_sequence_length,
+        FLAGS.tokenizer_vocab_size
+    )
+
+    encoder_model_dir = 'pretrained_visual_model'
+    encoder = CNN_Encoder(
+        encoder_model_dir,
+        FLAGS.visual_model_name,
+        FLAGS.visual_model_pop_layers,
+        FLAGS.encoder_layers,
+        FLAGS.tags_threshold,
+        num_tags=len(FLAGS.tags)
+    )
+
+    decoder = DistilGPT2.from_pretrained('distilgpt2')
+    optimizer = torch.optim.Adam(decoder.parameters(), lr=FLAGS.learning_rate)
+
+    encoder.to(device)
+    decoder.to(device)
+
+    checkpoint_path = hf_hub_download(
+        repo_id="TransformingBerry/CDGPT2_checkpoint",
+        filename="checkpoint.pth"
+    )
 
-encoder = CNN_Encoder('pretrained_visual_model', FLAGS.visual_model_name, FLAGS.visual_model_pop_layers,
-                      FLAGS.encoder_layers, FLAGS.tags_threshold, num_tags=len(FLAGS.tags))
-decoder = DistilGPT2.from_pretrained('distilgpt2')
+    if os.path.exists(checkpoint_path):
+        print(f"✅ Restoring from checkpoint: {checkpoint_path}")
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        encoder.load_state_dict(checkpoint['encoder_state_dict'])
+        decoder.load_state_dict(checkpoint['decoder_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("⚠️ No checkpoint found. Starting from scratch.")
 
-optimizer = torch.optim.Adam(decoder.parameters(), lr=FLAGS.learning_rate)
-# checkpoint_path = os.path.join(FLAGS.ckpt_path, "checkpoint.pth")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-encoder.to(device)
-decoder.to(device)
+    print("✅ Model initialized.")
 
-checkpoint_path = hf_hub_download(repo_id="TransformingBerry/CDGPT2_checkpoint", filename="checkpoint.pth")
 
 
-if os.path.exists(checkpoint_path):
-    print(f"Restoring from checkpoint: {checkpoint_path}")
-    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
-    encoder.load_state_dict(checkpoint['encoder_state_dict'])
-    decoder.load_state_dict(checkpoint['decoder_state_dict'])
-    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-else:
-    print("No checkpoint found. Starting from scratch.")
+init_model()
 
 def generate_report(image_bytes):
     image = Image.open(io.BytesIO(image_bytes))