Update app.py

app.py

@@ -17,23 +17,31 @@ app.config['temp_response'] = None
 app.config['generation_thread'] = None
 
 
-def initialize_model():
-    global pipe
-    try:
-      print('Downloading the model weights')
-      device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-      xm = load_model('transmitter', device=device)
-      model = load_model('text300M', device=device)
-      diffusion = diffusion_from_config(load_config('diffusion'))
-      return device, xm, model, diffusion
-    except Exception as e:
-      print(f"Error downloading the model: {e}")
-      return jsonify({"error": f"Failed to download model: {str(e)}"}), 500
+# def initialize_model():
+#     global pipe
+#     try:
+#       print('Downloading the model weights')
+#       device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+#       xm = load_model('transmitter', device=device)
+#       model = load_model('text300M', device=device)
+#       diffusion = diffusion_from_config(load_config('diffusion'))
+#       return device, xm, model, diffusion
+#     except Exception as e:
+#       print(f"Error downloading the model: {e}")
+#       return jsonify({"error": f"Failed to download model: {str(e)}"}), 500
 
 def generate_image_gif(prompt):
   global pipe
   if pipe is None:
-    device, xm, model, diffusion = initialize_model()
+    try:
+        print('Downloading the model weights')
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        xm = load_model('transmitter', device=device)
+        model = load_model('text300M', device=device)
+        diffusion = diffusion_from_config(load_config('diffusion'))
+        pipe = 'Model loaded'
+    except Exception as e:
+        print(f"Error downloading the model: {e}")
 
   try:
     batch_size = 1