Update app.py

app.py

@@ -23,18 +23,9 @@ torch.backends.cudnn.benchmark = False
 #torch.backends.cuda.preferred_linalg_library="cusolver"
 torch.set_float32_matmul_precision("highest")
 
-upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cuda'))
-
-if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-    text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
-        "stable_diffusion_3_medium", width=768, height=768, dtype="bfloat16"
-    )
-else:
-    text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
-        "stable_diffusion_3_medium", width=768, height=768, dtype="float32"
-    )
-    logging.warning("bfloat16 not supported, using float32")
-
+text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
+    "stable_diffusion_3_medium", device='cuda', width=768, height=768, dtype="bfloat16"
+)
 
 code = r'''
 import paramiko
@@ -57,6 +48,8 @@ def upload_to_ftp(filename):
         print(f"FTP upload error: {e}")
 '''
 
+upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cuda'))
+
 pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -84,7 +77,7 @@ def infer_30(
     sd_image.save(sd35_path,optimize=False,compress_level=0)
     pyx.upload_to_ftp(sd35_path)
     with torch.no_grad():
-        upscale2 = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
+        upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
     downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
     upscale_path = f"sd3keras_upscale_{timestamp}.png"
@@ -113,7 +106,7 @@ def infer_60(
     sd_image.save(sd35_path,optimize=False,compress_level=0)
     pyx.upload_to_ftp(sd35_path)
     with torch.no_grad():
-        upscale2 = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
+        upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
     downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
     upscale_path = f"sd3keras_upscale_{timestamp}.png"
@@ -142,7 +135,7 @@ def infer_90(
     sd_image.save(sd35_path,optimize=False,compress_level=0)
     pyx.upload_to_ftp(sd35_path)
     with torch.no_grad():
-        upscale2 = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
+        upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
     downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
     upscale_path = f"sd3keras_upscale_{timestamp}.png"