Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import subprocess
 subprocess.run(['sh', './spaces.sh'])
 import spaces
@@ -35,6 +34,10 @@ import random
 import datetime
 import threading
 import io
+from PIL import Image
+
+# For Ultra HDR
+import pillow_ultrahdr
 
 from google.oauth2 import service_account
 from google.cloud import storage
@@ -42,10 +45,15 @@ from google.cloud import storage
 import torch
 
 @spaces.GPU(required=True)
-def install_flashattn():
+def install_dependencies():
     subprocess.run(['sh', './flashattn.sh'])
+    # Install the UltraHDR library
+    print("Installing pillow-ultrahdr...")
+    subprocess.run(['pip', 'install', 'pillow-ultrahdr'])
+    print("✅ pillow-ultrahdr installed.")
 
-#install_flashattn()
+# Install all dependencies
+# install_dependencies()
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -58,7 +66,6 @@ torch.backends.cuda.preferred_linalg_library="cusolver"
 torch.set_float32_matmul_precision("highest")
 
 from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, AutoencoderKL
-from PIL import Image
 from image_gen_aux import UpscaleWithModel
 
 
@@ -73,23 +80,22 @@ if GCS_SA_KEY and GCS_BUCKET_NAME:
         gcs_client = storage.Client(credentials=credentials)
         print("✅ GCS Client initialized successfully.")
     except Exception as e:
-        print(f"❌ Failed to initialize GCS client: {e}")
+        print(f"❌ Failed to initialize GCS client: ")
 
-def upload_to_gcs(image_object, filename):
+def upload_to_gcs(image_bytes, filename):
     if not gcs_client:
         print("⚠️ GCS client not initialized. Skipping upload.")
         return
     try:
-        print(f"--> Starting GCS upload for {filename}...")
+        print(f"--> Starting GCS upload for {}...")
         bucket = gcs_client.bucket(GCS_BUCKET_NAME)
-        blob = bucket.blob(f"stablediff/{filename}")
-        img_byte_arr = io.BytesIO()
-        image_object.save(img_byte_arr, format='PNG', optimize=False, compress_level=0)
-        img_byte_arr = img_byte_arr.getvalue()
-        blob.upload_from_string(img_byte_arr, content_type='image/png')
-        print(f"✅ Successfully uploaded {filename} to GCS.")
+        blob = bucket.blob(f"stablediff/{}")
+        # The image_bytes is already a bytes object, so we can upload it directly
+        blob.upload_from_string(image_bytes, content_type='image/jpeg')
+        print(f"✅ Successfully uploaded {} to GCS.")
     except Exception as e:
-        print(f"❌ An error occurred during GCS upload: {e}")
+        print(f"❌ An error occurred during GCS upload for {}: {}")
+
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -136,7 +142,7 @@ class FlashAttentionProcessor(AttnProcessor2_0):
         out = out_reshaped.permute(0, 2, 1, 3).reshape(b, t, c)
         out = attn.to_out(out)
         return out
-        
+
 @spaces.GPU(duration=120)
 def compile_transformer():
     with spaces.aoti_capture(pipe.transformer) as call:
@@ -147,7 +153,7 @@ def compile_transformer():
         kwargs=call.kwargs,
     )
     return spaces.aoti_compile(exported)
-    
+
 def load_model():
     pipe = StableDiffusion3Pipeline.from_pretrained(
         "ford442/stable-diffusion-3.5-large-bf16",
@@ -162,6 +168,15 @@ def load_model():
     upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(device)
     return pipe, upscaler_2
 
+def srgb_to_linear(img_tensor):
+    """Converts a batched sRGB tensor [0, 1] to a linear tensor."""
+    # Using the standard sRGB to linear conversion formula
+    return torch.where(
+        img_tensor <= 0.04045,
+        img_tensor / 12.92,
+        ((img_tensor + 0.055) / 1.055).pow(2.4)
+    )
+
 pipe, upscaler_2 = load_model()
 
 fa_processor = FlashAttentionProcessor()
@@ -176,123 +191,94 @@ spaces.aoti_apply(compiled_transformer, pipe.transformer)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
-@spaces.GPU(duration=45)
-def generate_images_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress=gr.Progress(track_tqdm=True)):
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    print('-- generating image --')
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    sd_image = pipe(
-        prompt=prompt, prompt_2=prompt, prompt_3=prompt,
-        negative_prompt=neg_prompt_1, negative_prompt_2=neg_prompt_2, negative_prompt_3=neg_prompt_3,
-        guidance_scale=guidance, num_inference_steps=steps,
-        width=width, height=height, generator=generator,
-        max_sequence_length=384
-    ).images[0]
-    print('-- got image --')
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    with torch.no_grad():
-        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)
-    return sd_image, 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)):
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    print('-- generating image --')
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    sd_image = pipe(
-        prompt=prompt, prompt_2=prompt, prompt_3=prompt,
-        negative_prompt=neg_prompt_1, negative_prompt_2=neg_prompt_2, negative_prompt_3=neg_prompt_3,
-        guidance_scale=guidance, num_inference_steps=steps,
-        width=width, height=height, generator=generator,
-        max_sequence_length=384
-    ).images[0]
-    print('-- got image --')
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    with torch.no_grad():
-        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)
-    return sd_image, downscaled_upscale, prompt
+# Consolidated generation function
+def generate_images(duration, prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress=gr.Progress(track_tqdm=True)):
+    @spaces.GPU(duration=duration)
+    def _generate():
+        seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device=device).manual_seed(seed)
+        print('-- generating image --')
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+        
+        # Generate tensor output in sRGB space
+        sd_image_tensor_srgb = pipe(
+            prompt=prompt, prompt_2=prompt, prompt_3=prompt,
+            negative_prompt=neg_prompt_1, negative_prompt_2=neg_prompt_2, negative_prompt_3=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
+
+        # 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()
 
-@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)):
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    print('-- generating image --')
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    sd_image = pipe(
-        prompt=prompt, prompt_2=prompt, prompt_3=prompt,
-        negative_prompt=neg_prompt_1, negative_prompt_2=neg_prompt_2, negative_prompt_3=neg_prompt_3,
-        guidance_scale=guidance, num_inference_steps=steps,
-        width=width, height=height, generator=generator,
-        max_sequence_length=384
-    ).images[0]
-    print('-- got image --')
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    with torch.no_grad():
-        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)
-    return sd_image, 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)
-    if save_consent:
-        print("✅ User consented to save. Preparing uploads...")
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        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()
-    else:
-        print("ℹ️ User did not consent to save. Skipping upload.")
-    return sd_image, expanded_prompt
+        # 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
+    
+    return _generate()
 
-def run_inference_and_upload_60(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_60(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress)
+# Consolidated upload function
+def run_inference_and_upload(duration, prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, save_consent, progress=gr.Progress(track_tqdm=True)):
+    # Generate images and get both PIL (for display) and bytes (for upload)
+    sd_image_pil, sd_hdr_bytes, upscaled_hdr_bytes, expanded_prompt = generate_images(
+        duration, prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress
+    )
+    
     if save_consent:
         print("✅ User consented to save. Preparing uploads...")
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        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_filename = f"sd35ll_{}.jpg"
+        upscale_filename = f"sd35ll_upscale_{}.jpg"
+        
+        # Upload using threading
+        sd_thread = threading.Thread(target=upload_to_gcs, args=(sd_hdr_bytes, sd_filename))
+        upscale_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_hdr_bytes, upscale_filename))
         sd_thread.start()
         upscale_thread.start()
     else:
         print("ℹ️ User did not consent to save. Skipping upload.")
-    return sd_image, expanded_prompt
+        
+    # Return the standard sRGB PIL image to the Gradio interface for display
+    return sd_image_pil, expanded_prompt
 
-def run_inference_and_upload_110(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_110(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress)
-    if save_consent:
-        print("✅ User consented to save. Preparing uploads...")
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        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()
-    else:
-        print("ℹ️ User did not consent to save. Skipping upload.")
-    return sd_image, expanded_prompt
-    
 css = """
 #col-container {margin: 0 auto;max-width: 640px;}
 body{background-color: blue;}
@@ -309,11 +295,12 @@ with gr.Blocks(theme=gr.themes.Origin(), css=css) as demo:
             )
             run_button_30 = gr.Button("Run30", scale=0, variant="primary")
             run_button_60 = gr.Button("Run60", scale=0, variant="primary")
-            run_button_110 = gr.Button("Run100", scale=0, variant="primary")
-        result = gr.Image(label="Result", show_label=False, type="pil")
+            run_button_110 = gr.Button("Run110", scale=0, variant="primary")
+        # The result will display the standard PIL image, the HDR is saved/uploaded
+        result = gr.Image(label="Result (SDR Preview)", show_label=False, type="pil")
         save_consent_checkbox = gr.Checkbox(
-            label="✅ Anonymously upload result to a public gallery",
-            value=True, # Default to not uploading
+            label="✅ Anonymously upload result to a public gallery (as JPEG Ultra HDR)",
+            value=True,
             info="Check this box to help us by contributing your image."
         )
         with gr.Accordion("Advanced Settings", open=True):
@@ -327,50 +314,30 @@ with gr.Blocks(theme=gr.themes.Origin(), css=css) as demo:
                 guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=30.0, step=0.1, value=4.2)
                 num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=150, step=1, value=60)
 
+        # Clicks now call the same function with a different duration parameter
         run_button_30.click(
-            fn=run_inference_and_upload_30,
+            fn=lambda *args: run_inference_and_upload(45, *args),
             inputs=[
-                prompt,
-                negative_prompt_1,
-                negative_prompt_2,
-                negative_prompt_3,
-                width,
-                height,
-                guidance_scale,
-                num_inference_steps,
-                save_consent_checkbox # Pass the checkbox value
+                prompt, negative_prompt_1, negative_prompt_2, negative_prompt_3,
+                width, height, guidance_scale, num_inference_steps, save_consent_checkbox
             ],
             outputs=[result, expanded_prompt_output],
         )
 
         run_button_60.click(
-            fn=run_inference_and_upload_60,
+            fn=lambda *args: run_inference_and_upload(70, *args),
             inputs=[
-                prompt,
-                negative_prompt_1,
-                negative_prompt_2,
-                negative_prompt_3,
-                width,
-                height,
-                guidance_scale,
-                num_inference_steps,
-                save_consent_checkbox # Pass the checkbox value
+                prompt, negative_prompt_1, negative_prompt_2, negative_prompt_3,
+                width, height, guidance_scale, num_inference_steps, save_consent_checkbox
             ],
             outputs=[result, expanded_prompt_output],
         )
 
         run_button_110.click(
-            fn=run_inference_and_upload_110,
+            fn=lambda *args: run_inference_and_upload(120, *args),
             inputs=[
-                prompt,
-                negative_prompt_1,
-                negative_prompt_2,
-                negative_prompt_3,
-                width,
-                height,
-                guidance_scale,
-                num_inference_steps,
-                save_consent_checkbox # Pass the checkbox value
+                prompt, negative_prompt_1, negative_prompt_2, negative_prompt_3,
+                width, height, guidance_scale, num_inference_steps, save_consent_checkbox
             ],
             outputs=[result, expanded_prompt_output],
         )