Update app.py

app.py

@@ -105,23 +105,23 @@ def predict(prompt, guidance, steps, seed=1231231):
     if nsfw_content_detected:
         nsfw=gr.Button("🕹️NSFW🎨", scale=1)
 
-    central = pytz.timezone('US/Central')
-    safe_date_time = datetime.datetime.now().strftime("%Y%m%d")
-    replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
-    safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
-    filename = f"{safe_date_time}_{safe_prompt}.png"
-    
-
-    # Save the image
-    if len(results.images) > 0:
-        image_path = os.path.join("", filename)  # Specify your directory
-        results.images[0].save(image_path)
-        print(f"#Image saved as {image_path}")
-        encoded_image = encode_image(image)
-        html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
-        gr.Markdown(html_link)
-    
-
+    try: 
+        central = pytz.timezone('US/Central')
+        safe_date_time = datetime.datetime.now().strftime("%Y%m%d")
+        replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
+        safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
+        filename = f"{safe_date_time}_{safe_prompt}.png"
+        
+        # Save the image
+        if len(results.images) > 0:
+            image_path = os.path.join("", filename)  # Specify your directory
+            results.images[0].save(image_path)
+            print(f"#Image saved as {image_path}")
+            encoded_image = encode_image(image)
+            html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
+            #gr.Markdown(html_link)
+    except:
+        return results.images[0]
 
     return results.images[0] if len(results.images) > 0 else None