Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from PIL import Image
+import os
 
 # Diffusers
 from diffusers import (
@@ -20,6 +21,9 @@ import flax
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
 
+os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"]="false"
+os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]=".XX"
+os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"]="platform"
 
 def create_key(seed=0):
     return jax.random.PRNGKey(seed)
@@ -74,7 +78,10 @@ def infer(prompts, negative_prompts, image):
 
 
 title = "ControlNet + Stable Diffusion for Battlemaps"
-description = "Sketch your game battlemap and add some prompts to let the magic happen 🪄. Pretrained on battlemaps images. By Orgrim, Karm and Robin "
+description = """Sketch your game battlemap and add some prompts to let the magic happen 🪄.
+Pretrained on battlemaps images. 
+By Orgrim, Karm and Robin "
+"""
 # you need to pass inputs and outputs according to inference function
 gr.Interface(
     fn=infer,