Update app.py

app.py

@@ -298,38 +298,42 @@ def stable_diffusion_compose(prompt, scale, steps):
 
 
 def compose(prompt, version, guidance_scale, steps):
-    with th.no_grad():
-        if version == 'GLIDE':
-            clevr_model.to(cpu)
-            pipe.to(cpu)
-            model.to(device)
-            model_up.to(device)
-            return compose_language_descriptions(prompt, guidance_scale, steps)
-        elif version == 'Stable_Diffusion_1v_4':
-            clevr_model.to(cpu)
-            model.to(cpu)
-            model_up.to(cpu)
-            pipe.to(device)
-            return stable_diffusion_compose(prompt, guidance_scale, steps)
-        else:
-            pipe.to(cpu)
-            model.to(cpu)
-            model_up.to(cpu)
-            clevr_model.to(device)
-            # simple check
-            is_text = True
-            for char in prompt:
-                if char.isdigit():
-                    is_text = False
-                    break
-            if is_text:
-                img = Image.new('RGB', (512, 512), color=(255, 255, 255))
-                d = ImageDraw.Draw(img)
-                font = ImageFont.load_default()
-                d.text((0, 256), "input should be similar to the example using 2D coordinates.", fill=(0, 0, 0), font=font)
-                return img
+    try:
+        with th.no_grad():
+            if version == 'GLIDE':
+                clevr_model.to(cpu)
+                pipe.to(cpu)
+                model.to(device)
+                model_up.to(device)
+                return compose_language_descriptions(prompt, guidance_scale, steps)
+            elif version == 'Stable_Diffusion_1v_4':
+                clevr_model.to(cpu)
+                model.to(cpu)
+                model_up.to(cpu)
+                pipe.to(device)
+                return stable_diffusion_compose(prompt, guidance_scale, steps)
             else:
-                return compose_clevr_objects(prompt, guidance_scale, steps)
+                pipe.to(cpu)
+                model.to(cpu)
+                model_up.to(cpu)
+                clevr_model.to(device)
+                # simple check
+                is_text = True
+                for char in prompt:
+                    if char.isdigit():
+                        is_text = False
+                        break
+                if is_text:
+                    img = Image.new('RGB', (512, 512), color=(255, 255, 255))
+                    d = ImageDraw.Draw(img)
+                    font = ImageFont.load_default()
+                    d.text((0, 256), "input should be similar to the example using 2D coordinates.", fill=(0, 0, 0), font=font)
+                    return img
+                else:
+                    return compose_clevr_objects(prompt, guidance_scale, steps)
+    except Exception as e:
+        print(e)
+        return None
 
 
 examples_1 = 'a camel | a forest'