debug code

app.py

@@ -23,6 +23,9 @@ _model = InferenceUNetPseudo3D(
         hf_auth_token = os.environ.get('HUGGING_FACE_HUB_TOKEN', None)
 )
 
+if _model.failed == True:
+    exit()
+
 # gradio is illiterate. type hints make it go poopoo in pantsu.
 def generate(
         prompt = 'An elderly man having a great time in the park.',

makeavid_sd/inference.py

@@ -62,20 +62,32 @@ class InferenceUNetPseudo3D:
         self.hf_auth_token = hf_auth_token
 
         self.params: Dict[str, FrozenDict[str, Any]] = {}
-        unet, unet_params = UNetPseudo3DConditionModel.from_pretrained(
-                self.model_path,
-                subfolder = 'unet',
-                from_pt = False,
-                sample_size = (64, 64),
-                dtype = self.dtype,
-                param_dtype = dtypestr(self.dtype),
-                use_memory_efficient_attention = True,
-                use_auth_token = self.hf_auth_token
-        )
-        self.unet: UNetPseudo3DConditionModel = unet
-        unet_params = castto(self.dtype, self.unet, unet_params)
-        self.params['unet'] = FrozenDict(unet_params)
-        del unet_params
+        try:
+            import traceback
+            print('initializing unet')
+            unet, unet_params = UNetPseudo3DConditionModel.from_pretrained(
+                    self.model_path,
+                    subfolder = 'unet',
+                    from_pt = False,
+                    sample_size = (64, 64),
+                    dtype = self.dtype,
+                    param_dtype = dtypestr(self.dtype),
+                    use_memory_efficient_attention = True,
+                    use_auth_token = self.hf_auth_token
+            )
+            self.unet: UNetPseudo3DConditionModel = unet
+            print('casting unet params')
+            unet_params = castto(self.dtype, self.unet, unet_params)
+            print('storing unet params')
+            self.params['unet'] = FrozenDict(unet_params)
+            print('deleting unet params')
+            del unet_params
+        except Exception as e:
+            print(e)
+            traceback.print_exc()
+            self.failed = True
+            return
+        self.failed = False
         vae, vae_params = FlaxAutoencoderKL.from_pretrained(
                 self.model_path,
                 subfolder = 'vae',