Update app.py

app.py

@@ -9,7 +9,7 @@ from image_gen_aux import UpscaleWithModel
 import cyper
 from PIL import Image
 os.environ['JAX_PLATFORMS'] = 'cpu'
-
+import random
 os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"
 import keras
 import keras_hub
@@ -28,14 +28,15 @@ torch.backends.cudnn.benchmark = False
 torch.set_float32_matmul_precision("highest")
 
 upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(device)
+text_to_image = None
 
 def load_model():
+    global text_to_image
     text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
         "stable_diffusion_3_medium", width=768, height=768, dtype="bfloat16"
     )
     return text_to_image
 
-
 code = r'''
 import paramiko
 import os
@@ -57,7 +58,6 @@ def upload_to_ftp(filename):
         print(f"FTP upload error: {e}")
 '''
 
-
 pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -72,7 +72,9 @@ def infer_30(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    text_to_image = load_model()
+    global text_to_image
+    if text_to_image is None:
+        text_to_image = load_model()
     os.environ['JAX_PLATFORMS'] = 'gpu'
     os.environ['KERAS_BACKEND'] = 'jax'
     seed = random.randint(0, MAX_SEED)
@@ -104,7 +106,9 @@ def infer_60(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    text_to_image = load_model()
+    global text_to_image
+    if text_to_image is None:
+        text_to_image = load_model()
     os.environ['JAX_PLATFORMS'] = 'gpu'
     os.environ['KERAS_BACKEND'] = 'jax'
     seed = random.randint(0, MAX_SEED)
@@ -136,7 +140,9 @@ def infer_90(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    text_to_image = load_model()
+    global text_to_image
+    if text_to_image is None:
+        text_to_image = load_model()
     os.environ['JAX_PLATFORMS'] = 'gpu'
     os.environ['KERAS_BACKEND'] = 'jax'
     seed = random.randint(0, MAX_SEED)
@@ -178,6 +184,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
+            load_button = gr.Button("Load model", scale=0, variant="primary")
             run_button_30 = gr.Button("Run 30", scale=0, variant="primary")
             run_button_60 = gr.Button("Run 60", scale=0, variant="primary")
             run_button_90 = gr.Button("Run 90", scale=0, variant="primary")
@@ -206,6 +213,13 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                     value=50,
                 )
                 
+        gr.on(
+        triggers=[load_button.click],
+        fn=load_model,
+        inputs=[],
+        outputs=[],
+        )
+        
         gr.on(
         triggers=[run_button_30.click, prompt.submit],
         fn=infer_30,