Update app.py

app.py

@@ -8,9 +8,9 @@ import paramiko
 from image_gen_aux import UpscaleWithModel
 import cyper
 from PIL import Image
+os.environ['JAX_PLATFORMS'] = 'cpu'
 
 os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"
-os.environ['KERAS_BACKEND'] = 'jax'
 import keras
 import keras_hub
 import torch
@@ -73,6 +73,8 @@ def infer_30(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    os.environ['JAX_PLATFORMS'] = 'gpu'
+    os.environ['KERAS_BACKEND'] = 'jax'
     seed = random.randint(0, MAX_SEED)
     sd_image = text_to_image.generate(
         prompt=prompt,
@@ -102,6 +104,8 @@ def infer_60(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    os.environ['JAX_PLATFORMS'] = 'gpu'
+    os.environ['KERAS_BACKEND'] = 'jax'
     seed = random.randint(0, MAX_SEED)
     sd_image = text_to_image.generate(
         prompt=prompt,
@@ -131,6 +135,8 @@ def infer_90(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    os.environ['JAX_PLATFORMS'] = 'gpu'
+    os.environ['KERAS_BACKEND'] = 'jax'
     seed = random.randint(0, MAX_SEED)
     sd_image = text_to_image.generate(
         prompt=prompt,