Update app.py

app.py

@@ -18,6 +18,7 @@ from typing import Tuple
 import paramiko
 import datetime
 import cyper
+from image_gen_aux import UpscaleWithModel
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -81,6 +82,8 @@ os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
+
 def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         # adjust the batch_size of prompt_embeds according to guidance_scale
     if step_index == int(pipeline.num_timesteps * 0.1):
@@ -176,6 +179,9 @@ def uploadNote(prompt,num_inference_steps,guidance_scale,timestamp):
 
 code = r'''
 
+import torch
+import paramiko
+
 def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         # adjust the batch_size of prompt_embeds according to guidance_scale
     if step_index == int(pipeline.num_timesteps * 0.1):
@@ -241,7 +247,7 @@ def uploadNote(prompt,num_inference_steps,guidance_scale,timestamp):
     
 '''
 
-pyx = cyper.inline(code)
+pyx = cyper.inline(code, fast_indexing=True)
 
 @spaces.GPU(duration=30)
 def generate_30(
@@ -283,6 +289,13 @@ def generate_30(
     sd_image_path = f"rv_C_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     pyx.upload_to_ftp(sd_image_path)    
+    torch.set_float32_matmul_precision("medium")
+    with torch.no_grad():
+        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+    downscale_path = f"rv50_upscale_{timestamp}.png"
+    downscale1.save(downscale_path,optimize=False,compress_level=0)
+    pyx.upload_to_ftp(downscale_path) 
     unique_name = str(uuid.uuid4()) + ".png"  
     os.symlink(sd_image_path, unique_name)  
     return [unique_name]