Update app.py

app.py

@@ -1,8 +1,10 @@
-import spaces
 
 import subprocess
-#subprocess.run(['sh', './spaces.sh'])
-    
+
+
+subprocess.run(['sh', './spaces.sh'])
+import spaces
+
 @spaces.GPU(required=True)
 def install_dependencies():
     subprocess.run(['sh', './flashattn.sh'])
@@ -48,7 +50,6 @@ from PIL import Image
 # For  HDR
 import pillow_avif
 import cv2
-
 from google.oauth2 import service_account
 from google.cloud import storage
 
@@ -80,7 +81,7 @@ 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: ")
+        print(f"❌ Failed to initialize GCS client: {e}")
 
 def upload_to_gcs(image_bytes, filename):
     if not gcs_client:
@@ -90,13 +91,49 @@ def upload_to_gcs(image_bytes, filename):
         print(f"--> Starting GCS upload for {filename}...")
         bucket = gcs_client.bucket(GCS_BUCKET_NAME)
         blob = bucket.blob(f"stablediff/{filename}")
-        # 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 {filename} to GCS.")
+        blob.upload_from_string(image_bytes, content_type='image/avif')
+        print(f"✅ Successfully uploaded {} to GCS.")
     except Exception as e:
         print(f"❌ An error occurred during GCS upload for {filename}: {e}")
 
+def srgb_to_linear(tensor_srgb):
+    """Converts a batched sRGB PyTorch tensor [0, 1] to a linear tensor."""
+    return torch.where(
+        tensor_srgb <= 0.04045,
+        tensor_srgb / 12.92,
+        ((tensor_srgb + 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()
+    
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 from diffusers.models.attention_processor import AttnProcessor2_0
@@ -142,7 +179,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:
@@ -153,7 +190,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",
@@ -168,45 +205,6 @@ 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)
-    )
-
-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()
@@ -221,66 +219,126 @@ spaces.aoti_apply(compiled_transformer, pipe.transformer)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
-# 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="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, creating HDR AVIF version --')
-    sd_avif_bytes = create_hdr_avif_bytes(sd_image_srgb_pil)
-    
-    print('-- upscaling image --')
+@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_srgb_pil, tiling=True, tile_width=256, tile_height=256)
+        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, 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()
+    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
+
+@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)
+    sd_avif_bytes = create_hdr_avif_bytes(downscaled_upscale)
+    return sd_avif_bytes, sd_avif_bytes, prompt
 
-# 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
-    )
-    
+@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)
+    sd_avif_bytes = create_hdr_avif_bytes(downscaled_upscale)
+    return sd_avif_bytes, sd_avif_bytes, 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}.jpg"
-        upscale_filename = f"sd35ll_upscale_{timestamp}.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_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.start()
     else:
         print("ℹ️ User did not consent to save. Skipping upload.")
-        
-    # Return the standard sRGB PIL image to the Gradio interface for display
-    return sd_image_pil, expanded_prompt
+    return sd_image, expanded_prompt
+
+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)
+    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
 
+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;}
@@ -297,12 +355,11 @@ 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("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")
+            run_button_110 = gr.Button("Run100", scale=0, variant="primary")
+        result = gr.Image(label="Result", show_label=False, type="pil")
         save_consent_checkbox = gr.Checkbox(
-            label="✅ Anonymously upload result to a public gallery (as JPEG Ultra HDR)",
-            value=True,
+            label="✅ Anonymously upload result to a public gallery",
+            value=True, # Default to not uploading
             info="Check this box to help us by contributing your image."
         )
         with gr.Accordion("Advanced Settings", open=True):
@@ -316,30 +373,50 @@ 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=lambda *args: run_inference_and_upload(45, *args),
+            fn=run_inference_and_upload_30,
             inputs=[
-                prompt, negative_prompt_1, negative_prompt_2, negative_prompt_3,
-                width, height, guidance_scale, num_inference_steps, save_consent_checkbox
+                prompt,
+                negative_prompt_1,
+                negative_prompt_2,
+                negative_prompt_3,
+                width,
+                height,
+                guidance_scale,
+                num_inference_steps,
+                save_consent_checkbox # Pass the checkbox value
             ],
             outputs=[result, expanded_prompt_output],
         )
 
         run_button_60.click(
-            fn=lambda *args: run_inference_and_upload(70, *args),
+            fn=run_inference_and_upload_60,
             inputs=[
-                prompt, negative_prompt_1, negative_prompt_2, negative_prompt_3,
-                width, height, guidance_scale, num_inference_steps, save_consent_checkbox
+                prompt,
+                negative_prompt_1,
+                negative_prompt_2,
+                negative_prompt_3,
+                width,
+                height,
+                guidance_scale,
+                num_inference_steps,
+                save_consent_checkbox # Pass the checkbox value
             ],
             outputs=[result, expanded_prompt_output],
         )
 
         run_button_110.click(
-            fn=lambda *args: run_inference_and_upload(120, *args),
+            fn=run_inference_and_upload_110,
             inputs=[
-                prompt, negative_prompt_1, negative_prompt_2, negative_prompt_3,
-                width, height, guidance_scale, num_inference_steps, save_consent_checkbox
+                prompt,
+                negative_prompt_1,
+                negative_prompt_2,
+                negative_prompt_3,
+                width,
+                height,
+                guidance_scale,
+                num_inference_steps,
+                save_consent_checkbox # Pass the checkbox value
             ],
             outputs=[result, expanded_prompt_output],
         )