Update app.py

app.py

@@ -32,7 +32,7 @@ torch.set_float32_matmul_precision("highest")
 '''
 
 DESCRIPTIONXX = """
-    ## ⚡⚡⚡⚡ REALVISXL V5.0 BF16 (Tester D) ⚡⚡⚡⚡
+    ## ⚡⚡⚡⚡ REALVISXL V5.0 BF16 (Tester L) ⚡⚡⚡⚡
 """
 
 examples = [
@@ -273,43 +273,167 @@ def generate_30(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    denoising_start=0.0,
+    denoising_end =1.0,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    options = {
-        "prompt": [prompt],
-        "negative_prompt": [negative_prompt],
-        "negative_prompt_2": [neg_prompt_2],
-        "width": width,
-        "height": height,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "generator": generator,
-        "output_type": "pil",
-        "callback_on_step_end": pyx.scheduler_swap_callback
-    }
-    if use_resolution_binning:
-        options["use_resolution_binning"] = True
     images = []
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     filename = pyx.uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     #upload_to_ftp(filename) 
     pyx.upload_to_ftp(filename) 
     #uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
-    batch_options = options.copy()
-    rv_image = pipe(**batch_options).images[0]
-    sd_image_path = f"rv_D_{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"rv_D_upscale_{timestamp}.png"
-    downscale1.save(downscale_path,optimize=False,compress_level=0)
-    pyx.upload_to_ftp(downscale_path) 
+
+    if denoising_start=0.0:
+        options = {
+            "prompt": [prompt],
+            "negative_prompt": [negative_prompt],
+            "negative_prompt_2": [neg_prompt_2],
+            "width": width,
+            "height": height,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+            "generator": generator,
+            "output_type": "latent",
+            "denoising_start": denoising_start,
+            "denoising_end": denoising_end,
+            "callback_on_step_end": pyx.scheduler_swap_callback
+        }
+        if use_resolution_binning:
+            options["use_resolution_binning"] = True
+        batch_options = options.copy()
+        rv_image = pipe(**batch_options).images
+        sd_latent_path = f"rv_L_{denoising_end}.l"
+        torch.save(rv_image, sd_latent_path)
+    
+    if denoising_end!=1.0 and denoising_start!=0.0:
+        latent_file = f'rv_L_{denoising_start}.l'
+        loaded_latent = torch.load(latent_file)
+        print("Shape of loaded latent:", loaded_latent.shape)
+        # Get the expected shape for the latent tensor from the pipeline
+        batch_size = 1 # Usually 1 for single image generation
+        num_channels = pipe.vae.config.latent_channels  #typically 4
+        latent_shape = (batch_size, num_channels, height, width)
+        latents = torch.randn(latent_shape, dtype=torch.bfloat16, device=device)
+        if loaded_latent.shape == latents.shape:
+                print("Tensor already in expected shape. Not resizing...")
+                lat_file = loaded_latent.to(latents.device).to(latents.dtype)
+        elif loaded_latent.shape[1:] == latents.shape[1:]: #check if spatial dimensions are the same, in case of different batch sizes
+                latents[0] = loaded_latent[0].to(latents.device).to(latents.dtype) #copy only the first element if different batch size
+                lat_file = latents.to(latents.device).to(latents.dtype)
+        else:
+                print("Warning: Loaded latent shape does not match expected shape. Resizing...")
+                latents = torch.nn.functional.interpolate(loaded_latent, size=(height, width), mode='bilinear', align_corners=False).to(latents.device).to(latents.dtype)
+                if latents.shape[0] > 1:
+                    latents = latents[0].unsqueeze(0)
+                lat_file = latents.to(latents.device).to(latents.dtype)
+        print("-- using latent file --")
+        
+        options = {
+            "prompt": [prompt],
+            "negative_prompt": [negative_prompt],
+            "negative_prompt_2": [neg_prompt_2],
+            "width": width,
+            "height": height,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+            "generator": generator,
+            "latents": lat_file,
+            "output_type": "latent",
+            "denoising_start": denoising_start,
+            "denoising_end": denoising_end,
+            "callback_on_step_end": pyx.scheduler_swap_callback
+        }        
+        if use_resolution_binning:
+            options["use_resolution_binning"] = True
+        batch_options = options.copy()
+        rv_image = pipe(**batch_options).images
+        sd_latent_path = f"rv_L_{denoising_end}.l"
+        torch.save(rv_image, sd_latent_path)
+        
+    if denoising_end=1.0 and denoising_start!=0.0:
+        latent_file = f'rv_L_{denoising_start}.l'
+        loaded_latent = torch.load(latent_file)
+        print("Shape of loaded latent:", loaded_latent.shape)
+        # Get the expected shape for the latent tensor from the pipeline
+        batch_size = 1 # Usually 1 for single image generation
+        num_channels = pipe.vae.config.latent_channels  #typically 4
+        latent_shape = (batch_size, num_channels, height, width)
+        latents = torch.randn(latent_shape, dtype=torch.bfloat16, device=device)
+        if loaded_latent.shape == latents.shape:
+                print("Tensor already in expected shape. Not resizing...")
+                lat_file = loaded_latent.to(latents.device).to(latents.dtype)
+        elif loaded_latent.shape[1:] == latents.shape[1:]: #check if spatial dimensions are the same, in case of different batch sizes
+                latents[0] = loaded_latent[0].to(latents.device).to(latents.dtype) #copy only the first element if different batch size
+                lat_file = latents.to(latents.device).to(latents.dtype)
+        else:
+                print("Warning: Loaded latent shape does not match expected shape. Resizing...")
+                latents = torch.nn.functional.interpolate(loaded_latent, size=(height, width), mode='bilinear', align_corners=False).to(latents.device).to(latents.dtype)
+                if latents.shape[0] > 1:
+                    latents = latents[0].unsqueeze(0)
+                lat_file = latents.to(latents.device).to(latents.dtype)
+        print("-- using latent file --")
+        options = {
+            "prompt": [prompt],
+            "negative_prompt": [negative_prompt],
+            "negative_prompt_2": [neg_prompt_2],
+            "width": width,
+            "height": height,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+            "generator": generator,
+            "latents": lat_file,
+            "output_type": "pil",
+            "denoising_start": denoising_start,
+            "denoising_end": denoising_end,
+            "callback_on_step_end": pyx.scheduler_swap_callback
+        }       
+        if use_resolution_binning:
+            options["use_resolution_binning"] = True
+        batch_options = options.copy()
+        rv_image = pipe(**batch_options).images[0]
+        sd_image_path = f"rv_L_{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"rv_L_upscale_{timestamp}.png"
+        downscale1.save(downscale_path,optimize=False,compress_level=0)
+        pyx.upload_to_ftp(downscale_path) 
+
+    if denoising_end=1.0 and denoising_start=0.0:
+        options = {
+            "prompt": [prompt],
+            "negative_prompt": [negative_prompt],
+            "negative_prompt_2": [neg_prompt_2],
+            "width": width,
+            "height": height,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+            "generator": generator,
+            "output_type": "pil",
+            "callback_on_step_end": pyx.scheduler_swap_callback
+        }   
+       if use_resolution_binning:
+            options["use_resolution_binning"] = True
+        batch_options = options.copy()
+        rv_image = pipe(**batch_options).images[0]
+        sd_image_path = f"rv_L_{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"rv_L_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]
@@ -490,6 +614,22 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                 step=10,
                 value=170,
             )
+            denoising_start = gr.Slider(
+                label="Denoising Start",
+                minimum=0.0,
+                maximum=0.9,
+                step=0.1,
+                value=0.0,
+            )
+            
+            denoising_end = gr.Slider(
+                label="Denoising End",
+                minimum=0.1,
+                maximum=1.0,
+                step=0.1,
+                value=1.0,
+            )
+            
 
     gr.Examples(
         examples=examples,
@@ -519,6 +659,8 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            denoising_start,
+            denoising_end,
         ],
         outputs=[result],
     )