Update app.py

app.py

@@ -96,6 +96,26 @@ def upload_to_gcs(image_bytes, filename):
     except Exception as e:
         print(f"❌ An error occurred during GCS upload for {filename}: {e}")
 
+def srgb_to_linear_tensor(img_tensor_srgb):
+    """Converts a PyTorch sRGB tensor [0, 1] to a linear tensor."""
+    linear_mask = (img_tensor_srgb <= 0.04045).float()
+    non_linear_mask = (img_tensor_srgb > 0.04045).float()
+    linear_part = img_tensor_srgb / 12.92
+    non_linear_part = torch.pow((img_tensor_srgb + 0.055) / 1.055, 2.4)
+    img_linear = (linear_part * linear_mask) + (non_linear_part * non_linear_mask)
+    return img_linear
+
+def linear_to_srgb_tensor(img_tensor_linear):
+    """Converts a PyTorch linear tensor [0, 1] to sRGB."""
+    # Clamp to prevent negative values from torch.pow
+    img_tensor_linear = img_tensor_linear.clamp(min=0.0) 
+    srgb_mask = (img_tensor_linear <= 0.0031308).float()
+    non_srgb_mask = (img_tensor_linear > 0.0031308).float()
+    srgb_part = img_tensor_linear * 12.92
+    non_srgb_part = 1.055 * torch.pow(img_tensor_linear, 1.0/2.4) - 0.055
+    img_srgb = (srgb_part * srgb_mask) + (non_srgb_part * non_srgb_mask)
+    return img_srgb.clamp(0.0, 1.0)
+    
 def srgb_to_linear(tensor_srgb):
     """Converts a batched sRGB PyTorch tensor [0, 1] to a linear tensor."""
     return torch.where(
@@ -118,7 +138,7 @@ def create_hdr_avif_bytes(sdr_pil_image):
     hdr_16bit_array = (np.clip(linear_numpy_float, 0, 1) * 65535).astype(np.uint16)
 
     # 4. Create a PIL image that holds the 16-bit data
-    hdr_pil_image = Image.fromarray(hdr_16bit_array)
+    hdr_pil_image = Image.fromarray(hdr_16bit_array, mode='RGB;16')
 
     # 5. Save to a bytes buffer as 10-bit AVIF with HDR10 metadata
     buffer = io.BytesIO()
@@ -240,9 +260,9 @@ def generate_images_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width,
         upscale = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
         upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
-    downscaled_upscale = upscale2.resize((upscale2.width // 16, upscale2.height // 16), Image.LANCZOS)
-    sd_avif_bytes = create_hdr_avif_bytes(downscaled_upscale)
-    return sd_avif_bytes, sd_avif_bytes, prompt
+    sd_avif_bytes = create_hdr_avif_bytes(upscale2)
+    downscaled_upscale = sd_avif_bytes.resize((upscale2.width // 8, upscale2.height // 8), Image.LANCZOS)
+    return downscaled_upscale, downscaled_upscale, prompt
 
 @spaces.GPU(duration=70)
 def generate_images_60(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress=gr.Progress(track_tqdm=True)):
@@ -265,9 +285,9 @@ def generate_images_60(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width,
         upscale = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
         upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
-    downscaled_upscale = upscale2.resize((upscale2.width // 16, upscale2.height // 16), Image.LANCZOS)
-    sd_avif_bytes = create_hdr_avif_bytes(downscaled_upscale)
-    return sd_avif_bytes, sd_avif_bytes, prompt
+    sd_avif_bytes = create_hdr_avif_bytes(upscale2)
+    downscaled_upscale = sd_avif_bytes.resize((upscale2.width // 8, upscale2.height // 8), Image.LANCZOS)
+    return downscaled_upscale, downscaled_upscale, prompt
 
 @spaces.GPU(duration=120)
 def generate_images_110(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress=gr.Progress(track_tqdm=True)):
@@ -290,9 +310,9 @@ def generate_images_110(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width,
         upscale = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
         upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
-    downscaled_upscale = upscale2.resize((upscale2.width // 16, upscale2.height // 16), Image.LANCZOS)
-    sd_avif_bytes = create_hdr_avif_bytes(downscaled_upscale)
-    return sd_avif_bytes, sd_avif_bytes, prompt
+    sd_avif_bytes = create_hdr_avif_bytes(upscale2)
+    downscaled_upscale = sd_avif_bytes.resize((upscale2.width // 8, upscale2.height // 8), Image.LANCZOS)
+    return downscaled_upscale, downscaled_upscale, prompt
 
 def run_inference_and_upload_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, save_consent, progress=gr.Progress(track_tqdm=True)):
     sd_image, upscaled_image, expanded_prompt = generate_images_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress)
@@ -302,9 +322,9 @@ def run_inference_and_upload_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3
         sd_filename = f"sd35ll_{timestamp}.png"
         upscale_filename = f"sd35ll_upscale_{timestamp}.png"
         sd_thread = threading.Thread(target=upload_to_gcs, args=(sd_image, sd_filename))
-        #upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_image, upscale_filename))
-        sd_thread.start()
-        #upscale_thread.start()
+        upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_image, upscale_filename))
+        #sd_thread.start()
+        upscale_thread.start()
     else:
         print("ℹ️ User did not consent to save. Skipping upload.")
     return sd_image, expanded_prompt
@@ -317,9 +337,9 @@ def run_inference_and_upload_60(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3
         sd_filename = f"sd35ll_{timestamp}.png"
         upscale_filename = f"sd35ll_upscale_{timestamp}.png"
         sd_thread = threading.Thread(target=upload_to_gcs, args=(sd_image, sd_filename))
-        #upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_image, upscale_filename))
-        sd_thread.start()
-        #upscale_thread.start()
+        upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_image, upscale_filename))
+        #sd_thread.start()
+        upscale_thread.start()
     else:
         print("ℹ️ User did not consent to save. Skipping upload.")
     return sd_image, expanded_prompt
@@ -332,9 +352,9 @@ def run_inference_and_upload_110(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_
         sd_filename = f"sd35ll_{timestamp}.png"
         upscale_filename = f"sd35ll_upscale_{timestamp}.png"
         sd_thread = threading.Thread(target=upload_to_gcs, args=(sd_image, sd_filename))
-        #upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_image, upscale_filename))
-        sd_thread.start()
-        #upscale_thread.start()
+        upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_image, upscale_filename))
+        #sd_thread.start()
+        upscale_thread.start()
     else:
         print("ℹ️ User did not consent to save. Skipping upload.")
     return sd_image, expanded_prompt