import subprocess

subprocess.run(['sh', './spaces.sh'])
import spaces

@spaces.GPU(required=True)
def install_dependencies():
    subprocess.run(['sh', './flashattn.sh'])

install_dependencies()

import os

os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
os.environ['TORCH_LINALG_PREFER_CUSOLVER'] = '1'
os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True,pinned_use_background_threads:True'
os.environ["SAFETENSORS_FAST_GPU"] = "1"
os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1'

import torch

torch.backends.cuda.matmul.allow_tf32 = False  #  torch 2.8
torch.backends.cudnn.allow_tf32 = False        #  torch 2.8

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
#torch.backends.fp32_precision = "ieee"  torch 2.9
#torch.backends.cuda.matmul.fp32_precision = "ieee"  torch 2.9
#torch.backends.cudnn.fp32_precision = "ieee"  torch 2.9
#torch.backends.cudnn.conv.fp32_precision = "ieee"  torch 2.9
#torch.backends.cudnn.rnn.fp32_precision = "ieee"  torch 2.9
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False
torch.backends.cuda.preferred_blas_library="cublas"
torch.backends.cuda.preferred_linalg_library="cusolver"
torch.set_float32_matmul_precision("highest")
import json
import gradio as gr
import numpy as np
import random
import datetime
import threading
import io
from PIL import Image
import imageio.v3 as iio
# For  HDR
import pillow_avif
import cv2
from google.oauth2 import service_account
from google.cloud import storage

import torch

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False
torch.backends.cuda.preferred_blas_library="cublas"
torch.backends.cuda.preferred_linalg_library="cusolver"
torch.set_float32_matmul_precision("highest")

from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, AutoencoderKL
from image_gen_aux import UpscaleWithModel

GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
GCS_SA_KEY = os.getenv("GCS_SA_KEY") # The full JSON key content as a string
gcs_client = None

print(GCS_BUCKET_NAME)

if GCS_SA_KEY:
    print('Got key length: ',len(GCS_SA_KEY))
    
if GCS_SA_KEY and GCS_BUCKET_NAME:
    try:
        credentials_info = json.loads(GCS_SA_KEY) # New, safer way
        credentials = service_account.Credentials.from_service_account_info(credentials_info)
        gcs_client = storage.Client(credentials=credentials)
        print("✅ GCS Client initialized successfully.")
    except Exception as e:
        print(f"❌ Failed to initialize GCS client: {e}")

def upload_to_gcs(image_bytes, filename, content_type):
    if not gcs_client:
        print("⚠️ GCS client not initialized. Skipping upload.")
        return
    if image_bytes is None:
        print(f"⚠️ No image bytes for {filename}. Skipping upload.")
        return
    try:
        print(f"--> Starting GCS upload for {filename}...")
        bucket = gcs_client.bucket(GCS_BUCKET_NAME)
        blob = bucket.blob(f"stablediff/{filename}")
        
        # We already have the bytes, just upload them
        blob.upload_from_string(image_bytes, content_type=content_type)
        print(f"✅ Successfully uploaded {filename} to GCS.")
    except Exception as e:
        print(f"❌ An error occurred during GCS upload: {e}")

def srgb_to_linear_tensor(img_tensor_srgb):
    """Converts a PyTorch sRGB tensor [0, 1] to a linear tensor."""
    # Assumes input is in [0, 1] range
    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(
        tensor_srgb <= 0.04045,
        tensor_srgb / 12.92,
        ((tensor_srgb + 0.055) / 1.055).pow(2.4)
    )

def create_hdr_avif_bytes(image_tensor_fp32):
    """
    Converts a float32 sRGB tensor [-1, 1] to 10-bit HDR AVIF bytes.
    """
    if image_tensor_fp32 is None:
        return None
        
    try:
        # 1. Convert sRGB [-1, 1] tensor to Linear [0, 1] tensor
        srgb_tensor_0_1 = (image_tensor_fp32 / 2 + 0.5).clamp(0, 1)
        linear_tensor = srgb_to_linear_tensor(srgb_tensor_0_1)

        # 2. Convert linear float tensor to 16-bit uint numpy array [H, W, 3]
        # We use 16-bit as a container for our 10-bit data
        linear_tensor_16bit = (linear_tensor.clamp(0, 1) * 65535.0).round()
        hdr_16bit_array = linear_tensor_16bit.to(torch.uint16).cpu().permute(0, 2, 3, 1).numpy()[0]
        
        # 3. Save to bytes using imageio, forcing 10-bit YUV444 format for HDR
        return iio.imwrite(
            "<bytes>", 
            hdr_16bit_array, 
            format_hint=".avif", 
            codec="av1",         # Use AV1 codec
            out_pixelformat="yuv444p10le" # Force 10-bit, 4:4:4 chroma
        )
    except Exception as e:
        print(f"❌ Failed to encode HDR AVIF: {e}")
        return None
    
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

from diffusers.models.attention_processor import AttnProcessor2_0
from kernels import get_kernel
fa3_kernel = get_kernel("kernels-community/flash-attn3") # Or vllm-flash-attn3
class FlashAttentionProcessor(AttnProcessor2_0):
    def __call__(
        self,
        attn,
        hidden_states,
        encoder_hidden_states=None, # This will be present for cross-attention
        attention_mask=None,
        temb=None, # This might be present in some attention mechanisms, pass through if not used directly
        **kwargs,
    ):
        # Determine if it's self-attention or cross-attention
        # For self-attention, encoder_hidden_states is None or identical to hidden_states
        is_cross_attention = encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]
        # SD3.5 uses DiT, where hidden_states are often 3D (B, Seq, Dim)
        # However, attention can be within a transformer block which might internally reshape.
        # Ensure your inputs (query, key, value) are properly shaped for the kernel.
        # The kernel expects (Batch, Heads, Sequence, Dim_Head)
        query = attn.to_q(hidden_states)
        if is_cross_attention:
            key = attn.to_k(encoder_hidden_states)
            value = attn.to_v(encoder_hidden_states)
        else: # Self-attention
            key = attn.to_k(hidden_states)
            value = attn.to_v(hidden_states)
        scale = attn.scale
        query = query * scale
        b, t, c = query.shape # B=batch_size, T=sequence_length, C=embedding_dim
        h = attn.heads
        d = c // h # dim_per_head
        # Reshape to (Batch, Heads, Sequence, Dim_Head) for Flash Attention kernel
        q_reshaped = query.reshape(b, t, h, d).permute(0, 2, 1, 3)
        k_reshaped = key.reshape(b, t, h, d).permute(0, 2, 1, 3)
        v_reshaped = value.reshape(b, t, h, d).permute(0, 2, 1, 3)
        out_reshaped = torch.empty_like(q_reshaped)
        # Call the Flash Attention kernel
        fa3_kernel.attention(q_reshaped, k_reshaped, v_reshaped, out_reshaped)
        # Reshape output back to (Batch, Sequence, Heads * Dim_Head)
        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:
        pipe("A majestic, ancient Egyptian Sphinx stands sentinel in a large, clear pool under a bright, golden desert sun. Around its weathered stone base, several sleek, playful dolphins gracefully navigate the turquoise waters. The surrounding environment features lush, exotic papyrus plants and distant pyramids under a cloudless sky, conveying a sense of timeless wonder and serene majesty.")
    exported = torch.export.export(
        pipe.transformer,
        args=call.args,
        kwargs=call.kwargs,
    )
    return spaces.aoti_compile(exported)
    
def load_model():
    vae = AutoencoderKL.from_pretrained(
        "ford442/stable-diffusion-3.5-large-bf16", 
        subfolder="vae",
        torch_dtype=torch.float32  # Load VAE in full precision
    )
    pipe = StableDiffusion3Pipeline.from_pretrained(
        "ford442/stable-diffusion-3.5-large-bf16",
        trust_remote_code=True,
        transformer=None, # Load transformer separately
        use_safetensors=True,
        #vae=vae
    )
    ll_transformer=SD3Transformer2DModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='transformer').to(device, dtype=torch.bfloat16)
    pipe.transformer=ll_transformer
    pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")
    pipe.to(device=device, dtype=torch.bfloat16)
    pipe.vae=vae.to(device=device)
    upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(device)
    return pipe, upscaler_2

pipe, upscaler_2 = load_model()

fa_processor = FlashAttentionProcessor()

for name, module in pipe.transformer.named_modules():
    if isinstance(module, AttnProcessor2_0):
        module.processor = fa_processor

#compiled_transformer = compile_transformer()
#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,
        output_type="latent"  # <-- Get latents, not an image
        ).images
    
    # 2. Manually decode with our float32 VAE
    latents_fp32 = sd_image.to(torch.float32)
    latents_fp32 = 1 / pipe.vae.config.scaling_factor * latents_fp32
    with torch.no_grad():
        # This is our high-precision sRGB tensor in range [-1, 1]
        image_tensor_fp32 = pipe.vae.decode(latents_fp32).sample
    
    print('-- got fp32 image tensor --')
    
    # 3. Create 8-bit PIL Image from the tensor (for Gradio display)
    srgb_tensor_0_1 = (image_tensor_fp32 / 2 + 0.5).clamp(0, 1)
    srgb_numpy_8bit = (srgb_tensor_0_1.cpu().permute(0, 2, 3, 1).float().numpy()[0] * 255).round().astype("uint8")
    sd_image_pil_8bit = Image.fromarray(srgb_numpy_8bit)
    print('-- got 8-bit PIL image for display --')

    # 4. Create 16-bit PIL Image from the same tensor (for Upscaler)
    # Pillow 10+ can create an 'RGB' mode image from a uint16 array
    srgb_numpy_16bit = (srgb_tensor_0_1.cpu().permute(0, 2, 3, 1).float().numpy()[0] * 65535.0).round().astype("uint16")
    sd_image_pil_16bit = Image.fromarray(srgb_numpy_16bit)
    print('-- got 16-bit PIL image for upscaling --')

    # 5. Run the 16-bit upscaling (4x)
    # We feed the high-precision 16-bit PIL image to the upscaler
    with torch.no_grad():
        upscale_1 = upscaler_2(sd_image_pil_16bit, tiling=True, tile_width=256, tile_height=256)
        upscale_2 = upscaler_2(upscale_1, tiling=True, tile_width=256, tile_height=256)
    print('-- got 4K 16-bit upscaled PIL image --')

    torch.cuda.empty_cache()
    
    # 6. Convert the 4K 16-bit PIL back to a float32 tensor
    upscaled_16bit_numpy = np.array(upscale_2)
    upscaled_srgb_tensor = torch.from_numpy(upscaled_16bit_numpy).permute(2, 0, 1).unsqueeze(0).to(device, dtype=torch.float32) / 65535.0
    
    # 7. Create 10-bit HDR AVIF bytes from the 4K tensor (for GCS)
    # We pass the upscaled sRGB tensor, which create_hdr_avif_bytes will convert to linear
    # Note: We must convert the tensor from [0, 1] range back to [-1, 1] for create_hdr_avif_bytes
    upscaled_tensor_neg1_1 = (upscaled_srgb_tensor * 2.0 - 1.0).clamp(-1, 1)
    upscaled_avif_bytes = create_hdr_avif_bytes(upscaled_tensor_neg1_1)
    print('-- got 4K HDR AVIF bytes for upload --')

    # 8. Return the 8-bit PIL (for display) and the 4K AVIF bytes (for upload)
    return sd_image_pil_8bit, upscaled_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,
        output_type="latent"  # <-- Get latents, not an image
        ).images
    
    # 2. Manually decode with our float32 VAE
    latents_fp32 = sd_image.to(torch.float32)
    latents_fp32 = 1 / pipe.vae.config.scaling_factor * latents_fp32
    with torch.no_grad():
        # This is our high-precision sRGB tensor in range [-1, 1]
        image_tensor_fp32 = pipe.vae.decode(latents_fp32).sample
    
    print('-- got fp32 image tensor --')
    
    # 3. Create 8-bit PIL Image from the tensor (for Gradio display)
    srgb_tensor_0_1 = (image_tensor_fp32 / 2 + 0.5).clamp(0, 1)
    srgb_numpy_8bit = (srgb_tensor_0_1.cpu().permute(0, 2, 3, 1).float().numpy()[0] * 255).round().astype("uint8")
    sd_image_pil_8bit = Image.fromarray(srgb_numpy_8bit)
    print('-- got 8-bit PIL image for display --')

    # 4. Create 16-bit PIL Image from the same tensor (for Upscaler)
    # Pillow 10+ can create an 'RGB' mode image from a uint16 array
    srgb_numpy_16bit = (srgb_tensor_0_1.cpu().permute(0, 2, 3, 1).float().numpy()[0] * 65535.0).round().astype("uint16")
    sd_image_pil_16bit = Image.fromarray(srgb_numpy_16bit, mode='RGB')
    print('-- got 16-bit PIL image for upscaling --')

    # 5. Run the 16-bit upscaling (4x)
    # We feed the high-precision 16-bit PIL image to the upscaler
    with torch.no_grad():
        upscale_1 = upscaler_2(sd_image_pil_16bit, tiling=True, tile_width=256, tile_height=256)
        upscale_2 = upscaler_2(upscale_1, tiling=True, tile_width=256, tile_height=256)
    print('-- got 4K 16-bit upscaled PIL image --')

    torch.cuda.empty_cache()
    
    # 6. Convert the 4K 16-bit PIL back to a float32 tensor
    upscaled_16bit_numpy = np.array(upscale_2)
    upscaled_srgb_tensor = torch.from_numpy(upscaled_16bit_numpy).permute(2, 0, 1).unsqueeze(0).to(device, dtype=torch.float32) / 65535.0
    
    # 7. Create 10-bit HDR AVIF bytes from the 4K tensor (for GCS)
    # We pass the upscaled sRGB tensor, which create_hdr_avif_bytes will convert to linear
    # Note: We must convert the tensor from [0, 1] range back to [-1, 1] for create_hdr_avif_bytes
    upscaled_tensor_neg1_1 = (upscaled_srgb_tensor * 2.0 - 1.0).clamp(-1, 1)
    upscaled_avif_bytes = create_hdr_avif_bytes(upscaled_tensor_neg1_1)
    print('-- got 4K HDR AVIF bytes for upload --')

    # 8. Return the 8-bit PIL (for display) and the 4K AVIF bytes (for upload)
    return sd_image_pil_8bit, upscaled_avif_bytes, 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)):
    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,
        output_type="latent"  # <-- Get latents, not an image
        ).images
    
    # 2. Manually decode with our float32 VAE
    latents_fp32 = sd_image.to(torch.float32)
    latents_fp32 = 1 / pipe.vae.config.scaling_factor * latents_fp32
    with torch.no_grad():
        # This is our high-precision sRGB tensor in range [-1, 1]
        image_tensor_fp32 = pipe.vae.decode(latents_fp32).sample
    
    print('-- got fp32 image tensor --')
    
    # 3. Create 8-bit PIL Image from the tensor (for Gradio display)
    srgb_tensor_0_1 = (image_tensor_fp32 / 2 + 0.5).clamp(0, 1)
    srgb_numpy_8bit = (srgb_tensor_0_1.cpu().permute(0, 2, 3, 1).float().numpy()[0] * 255).round().astype("uint8")
    sd_image_pil_8bit = Image.fromarray(srgb_numpy_8bit)
    print('-- got 8-bit PIL image for display --')

    # 4. Create 16-bit PIL Image from the same tensor (for Upscaler)
    # Pillow 10+ can create an 'RGB' mode image from a uint16 array
    srgb_numpy_16bit = (srgb_tensor_0_1.cpu().permute(0, 2, 3, 1).float().numpy()[0] * 65535.0).round().astype("uint16")
    sd_image_pil_16bit = Image.fromarray(srgb_numpy_16bit, mode='RGB')
    print('-- got 16-bit PIL image for upscaling --')

    # 5. Run the 16-bit upscaling (4x)
    # We feed the high-precision 16-bit PIL image to the upscaler
    with torch.no_grad():
        upscale_1 = upscaler_2(sd_image_pil_16bit, tiling=True, tile_width=256, tile_height=256)
        upscale_2 = upscaler_2(upscale_1, tiling=True, tile_width=256, tile_height=256)
    print('-- got 4K 16-bit upscaled PIL image --')

    torch.cuda.empty_cache()
    
    # 6. Convert the 4K 16-bit PIL back to a float32 tensor
    upscaled_16bit_numpy = np.array(upscale_2)
    upscaled_srgb_tensor = torch.from_numpy(upscaled_16bit_numpy).permute(2, 0, 1).unsqueeze(0).to(device, dtype=torch.float32) / 65535.0
    
    # 7. Create 10-bit HDR AVIF bytes from the 4K tensor (for GCS)
    # We pass the upscaled sRGB tensor, which create_hdr_avif_bytes will convert to linear
    # Note: We must convert the tensor from [0, 1] range back to [-1, 1] for create_hdr_avif_bytes
    upscaled_tensor_neg1_1 = (upscaled_srgb_tensor * 2.0 - 1.0).clamp(-1, 1)
    upscaled_avif_bytes = create_hdr_avif_bytes(upscaled_tensor_neg1_1)
    print('-- got 4K HDR AVIF bytes for upload --')

    # 8. Return the 8-bit PIL (for display) and the 4K AVIF bytes (for upload)
    return sd_image_pil_8bit, upscaled_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)):
    
    # 1. Get the 8-bit PIL (for display) and 4K AVIF bytes (for upload)
    sd_image_pil, upscaled_avif_bytes, 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")
        
        # Define filenames
        sd_filename_png = f"sd35ll_{timestamp}.png"
        sd_filename_avif = f"sd35ll_4K_hdr_{timestamp}.avif"
        
        # 2. Convert the 8-bit PIL image to PNG bytes for upload
        img_byte_arr = io.BytesIO()
        sd_image_pil.save(img_byte_arr, format='PNG', optimize=False, compress_level=0)
        sd_png_bytes = img_byte_arr.getvalue()

        # 3. Start threads to upload both
        # Upload the 1K 8-bit PNG
        png_thread = threading.Thread(target=upload_to_gcs, args=(sd_png_bytes, sd_filename_png, "image/png"))
        # Upload the 4K 10-bit HDR AVIF
        avif_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_avif_bytes, sd_filename_avif, "image/avif"))
        
        png_thread.start()
        avif_thread.start()
    else:
        print("ℹ️ User did not consent to save. Skipping upload.")
        
    # 4. Return the 8-bit PIL image to Gradio
    return sd_image_pil, 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)):
    
    # 1. Get the 8-bit PIL (for display) and 4K AVIF bytes (for upload)
    sd_image_pil, upscaled_avif_bytes, 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")
        
        # Define filenames
        sd_filename_png = f"sd35ll_{timestamp}.png"
        sd_filename_avif = f"sd35ll_4K_hdr_{timestamp}.avif"
        
        # 2. Convert the 8-bit PIL image to PNG bytes for upload
        img_byte_arr = io.BytesIO()
        sd_image_pil.save(img_byte_arr, format='PNG', optimize=False, compress_level=0)
        sd_png_bytes = img_byte_arr.getvalue()

        # 3. Start threads to upload both
        # Upload the 1K 8-bit PNG
        png_thread = threading.Thread(target=upload_to_gcs, args=(sd_png_bytes, sd_filename_png, "image/png"))
        # Upload the 4K 10-bit HDR AVIF
        avif_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_avif_bytes, sd_filename_avif, "image/avif"))
        
        png_thread.start()
        avif_thread.start()
    else:
        print("ℹ️ User did not consent to save. Skipping upload.")
        
    # 4. Return the 8-bit PIL image to Gradio
    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)):
    
    # 1. Get the 8-bit PIL (for display) and 4K AVIF bytes (for upload)
    sd_image_pil, upscaled_avif_bytes, 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")
        
        # Define filenames
        sd_filename_png = f"sd35ll_{timestamp}.png"
        sd_filename_avif = f"sd35ll_4K_hdr_{timestamp}.avif"
        
        # 2. Convert the 8-bit PIL image to PNG bytes for upload
        img_byte_arr = io.BytesIO()
        sd_image_pil.save(img_byte_arr, format='PNG', optimize=False, compress_level=0)
        sd_png_bytes = img_byte_arr.getvalue()

        # 3. Start threads to upload both
        # Upload the 1K 8-bit PNG
        png_thread = threading.Thread(target=upload_to_gcs, args=(sd_png_bytes, sd_filename_png, "image/png"))
        # Upload the 4K 10-bit HDR AVIF
        avif_thread = threading.Thread(target=upload_to_gcs, args=(upscaled_avif_bytes, sd_filename_avif, "image/avif"))
        
        png_thread.start()
        avif_thread.start()
    else:
        print("ℹ️ User did not consent to save. Skipping upload.")
        
    # 4. Return the 8-bit PIL image to Gradio
    return sd_image_pil, expanded_prompt
    
css = """
#col-container {margin: 0 auto;max-width: 640px;}
body{background-color: blue;}
"""

with gr.Blocks(theme=gr.themes.Origin(), css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # StableDiffusion 3.5 Large with UltraReal lora test")
        expanded_prompt_output = gr.Textbox(label="Prompt", lines=1)
        with gr.Row():
            prompt = gr.Text(
                label="Prompt", show_label=False, max_lines=1,
                placeholder="Enter your prompt", container=False,
            )
            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")
        save_consent_checkbox = gr.Checkbox(
            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):
            negative_prompt_1 = gr.Text(label="Negative prompt 1", max_lines=1, placeholder="Enter a negative prompt", value="bad anatomy, poorly drawn hands, distorted face, blurry, out of frame, low resolution, grainy, pixelated, disfigured, mutated, extra limbs, bad composition")
            negative_prompt_2 = gr.Text(label="Negative prompt 2", max_lines=1, placeholder="Enter a second negative prompt", value="unrealistic, cartoon, anime, sketch, painting, drawing, illustration, graphic, digital art, render, 3d, blurry, deformed, disfigured, poorly drawn, bad anatomy, mutated, extra limbs, ugly, out of frame, bad composition, low resolution, grainy, pixelated, noisy, oversaturated, undersaturated, (worst quality, low quality:1.3), (bad hands, missing fingers:1.2)")
            negative_prompt_3 = gr.Text(label="Negative prompt 3", max_lines=1, placeholder="Enter a third negative prompt", value="(worst quality, low quality:1.3), (bad anatomy, bad hands, missing fingers, extra digit, fewer digits:1.2), (blurry:1.1), cropped, watermark, text, signature, logo, jpeg artifacts, (ugly, deformed, disfigured:1.2), (poorly drawn:1.2), mutated, extra limbs, (bad proportions, gross proportions:1.2), (malformed limbs, missing arms, missing legs, extra arms, extra legs:1.2), (fused fingers, too many fingers, long neck:1.2), (unnatural body, unnatural pose:1.1), out of frame, (bad composition, poorly composed:1.1), (oversaturated, undersaturated:1.1), (grainy, pixelated:1.1), (low resolution, noisy:1.1), (unrealistic, distorted:1.1), (extra fingers, mutated hands, poorly drawn hands, bad hands:1.3), (missing fingers:1.3)")
            with gr.Row():
                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            with gr.Row():
                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)

        run_button_30.click(
            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 # Pass the checkbox value
            ],
            outputs=[result, expanded_prompt_output],
        )

        run_button_60.click(
            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 # Pass the checkbox value
            ],
            outputs=[result, expanded_prompt_output],
        )

        run_button_110.click(
            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 # Pass the checkbox value
            ],
            outputs=[result, expanded_prompt_output],
        )

if __name__ == "__main__":
    demo.launch()