Update app.py

app.py

@@ -45,8 +45,9 @@ import threading
 import io
 from PIL import Image
 
-# For Ultra HDR
-import pillow_ultrahdr
+# For  HDR
+import pillow_avif
+import cv2
 
 from google.oauth2 import service_account
 from google.cloud import storage
@@ -176,6 +177,36 @@ def srgb_to_linear(img_tensor):
         ((img_tensor + 0.055) / 1.055).pow(2.4)
     )
 
+def create_hdr_avif_bytes(sdr_pil_image):
+    """Converts an SDR PIL image to a 10-bit HDR AVIF byte buffer."""
+    # 1. Convert SDR PIL image to a float tensor [0, 1]
+    srgb_tensor = torch.from_numpy(np.array(sdr_pil_image)).float().to(device) / 255.0
+    srgb_tensor = srgb_tensor.permute(2, 0, 1).unsqueeze(0)  # HWC to BCHW
+
+    # 2. Convert sRGB tensor to linear space
+    linear_tensor = srgb_to_linear(srgb_tensor)
+
+    # 3. Convert to 16-bit NumPy array for high-bit-depth processing
+    linear_numpy_float = linear_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
+    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)
+
+    # 5. Save to a bytes buffer as 10-bit AVIF with HDR10 metadata
+    buffer = io.BytesIO()
+    hdr_pil_image.save(
+        buffer,
+        format="AVIF",
+        quality=90,
+        depth=10,  # Specify 10-bit depth
+        subsampling="4:4:4",
+        color_primaries=9,  # BT.2020
+        transfer_characteristics=16,  # PQ (Perceptual Quantizer)
+        matrix_coefficients=9,  # BT.2020 non-constant luminance
+    )
+    return buffer.getvalue()
+    
 pipe, upscaler_2 = load_model()
 
 fa_processor = FlashAttentionProcessor()
@@ -207,50 +238,22 @@ def generate_images(duration, prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3,
             guidance_scale=guidance, num_inference_steps=steps,
             width=width, height=height, generator=generator,
             max_sequence_length=384,
-            output_type="pt"  # Request tensor output
-        ).images
+            output_type="pil"  # Get PIL for display and easy upscaling
+            ).images[0]
 
         # Convert the sRGB tensor [0,1] to a PIL Image for display and upscaling
         sd_image_pil_srgb = Image.fromarray((sd_image_tensor_srgb.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
-        print('-- got image --')
-        
-        # --- Upscaling ---
-        torch.cuda.empty_cache()
-        torch.cuda.reset_peak_memory_stats()
-        with torch.no_grad():
-            upscale = upscaler_2(sd_image_pil_srgb, tiling=True, tile_width=256, tile_height=256)
-            upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
-        print('-- got upscaled image --')
-
-        # --- HDR Conversion and Saving ---
-        # Convert the original sRGB tensor to linear space
-        sd_image_tensor_linear = srgb_to_linear(sd_image_tensor_srgb)
-        # Convert the linear tensor to a PIL Image (this will be HDR data)
-        sd_image_pil_linear = Image.fromarray((sd_image_tensor_linear.squeeze(0).permute(1, 2, 0).clamp(0, 1).cpu().numpy() * 255).astype(np.uint8))
-
-        # Save to a bytes buffer as JPEG Ultra HDR
-        buffer = io.BytesIO()
-        pillow_ultrahdr.save_ultrahdr(
-            sdr=sd_image_pil_srgb, # The standard dynamic range image
-            hdr=sd_image_pil_linear, # The linear (high dynamic range) image
-            outfile=buffer,
-            quality=90 # Standard JPEG quality setting
-        )
-        hdr_image_bytes = buffer.getvalue()
-        
-        # For the upscaled image, we will do the same
-        # First convert upscaled PIL image to tensor, normalize to [0,1]
-        upscaled_tensor_srgb = torch.from_numpy(np.array(upscale2)).float().to(device) / 255.0
-        upscaled_tensor_srgb = upscaled_tensor_srgb.permute(2, 0, 1).unsqueeze(0) # HWC to BCHW
-        upscaled_tensor_linear = srgb_to_linear(upscaled_tensor_srgb)
-        upscaled_pil_linear = Image.fromarray((upscaled_tensor_linear.squeeze(0).permute(1, 2, 0).clamp(0, 1).cpu().numpy() * 255).astype(np.uint8))
-        
-        upscaled_buffer = io.BytesIO()
-        pillow_ultrahdr.save_ultrahdr(sdr=upscale2, hdr=upscaled_pil_linear, outfile=upscaled_buffer, quality=95)
-        upscaled_hdr_image_bytes = upscaled_buffer.getvalue()
-
-        # Return the sRGB PIL image for display, and the HDR bytes for upload
-        return sd_image_pil_srgb, hdr_image_bytes, upscaled_hdr_image_bytes, prompt
+        print('-- got image, creating HDR AVIF version --')
+    sd_avif_bytes = create_hdr_avif_bytes(sd_image_srgb_pil)
+    
+    print('-- upscaling image --')
+    with torch.no_grad():
+        upscale = upscaler_2(sd_image_srgb_pil, tiling=True, tile_width=256, tile_height=256)
+        upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
+    print('-- got upscaled image, creating upscaled HDR AVIF --')
+    upscaled_avif_bytes = create_hdr_avif_bytes(upscale2)
+    
+    return sd_image_srgb_pil, sd_avif_bytes, upscaled_avif_bytes, prompt
     
     return _generate()