Update app.py

app.py

@@ -2,429 +2,206 @@ import subprocess
 subprocess.run(['sh', './spaces.sh'])
 
 import os
-
+# Environment variable setup
 os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
 os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
-alloc_conf_parts = [
-    'expandable_segments:True',
-    'pinned_use_background_threads:True'  # Specific to pinned memory.
-]
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(alloc_conf_parts)
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,pinned_use_background_threads:True'
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
 os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
 
 import spaces
-
 import gradio as gr
 import numpy as np
 import random
+import datetime
+import threading
+import io
+
+# --- New GCS Imports ---
+from google.oauth2 import service_account
+from google.cloud import storage
 
 import torch
+# Torch performance settings
 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 transformers import CLIPTextModelWithProjection, T5EncoderModel
-from transformers import CLIPTokenizer, T5TokenizerFast
-
-import re
-import paramiko
-import urllib
-import time
-from image_gen_aux import UpscaleWithModel
-from huggingface_hub import hf_hub_download
-import datetime
 from PIL import Image
+from image_gen_aux import UpscaleWithModel
 
-#from accelerate import Accelerator
-
-#accelerator = Accelerator(mixed_precision="bf16")
-
-hftoken = os.getenv("HF_AUTH_TOKEN") 
-
-import paramiko
-import socket
-import threading # NEW IMPORT
-import queue     # NEW IMPORT
-
-FTP_HOST = 'noahcohn.com'
-FTP_USER = 'ford442'
-FTP_PASS = os.getenv("FTP_PASS")
-FTP_DIR = 'img.noahcohn.com/stablediff/'
-FTP_HOST_FALLBACK = '1ink.us'
-FTP_DIR_FALLBACK = 'img.1ink.us/stablediff/'
+# --- GCS Configuration ---
+# Make sure to set these secrets in your Hugging Face Space settings
+GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
+GCS_SA_KEY = os.getenv("GCS_SA_KEY") # The full JSON key content as a string
 
-# --- WORKER FUNCTION FOR THREADING ---
-# This function contains the logic to connect to a single host.
-# It will be executed by each of our threads.
-def connect_worker(host, result_queue):
-    """Tries to connect to a single host and puts the successful transport object into the queue."""
-    transport = None
+# Initialize GCS client if credentials are available
+gcs_client = None
+if GCS_SA_KEY and GCS_BUCKET_NAME:
     try:
-        transport = paramiko.Transport((host, 22))
-        # We still use the 5-second timeout for the handshake
-        transport.start_client(timeout=5)
-        transport.auth_password(username=FTP_USER, password=FTP_PASS)
-        
-        # If we reach here, the connection was successful.
-        # Put the result in the queue for the main thread to use.
-        print(f"✅ Connection to {host} succeeded first.")
-        result_queue.put(transport)
-    except (paramiko.SSHException, socket.timeout, EOFError) as e:
-        # This is an expected failure, just print a note.
-        print(f"ℹ️ Connection to {host} failed or was too slow: {e}")
-        if transport:
-            transport.close()
+        credentials_info = eval(GCS_SA_KEY) # Using eval is safe here if you trust the secret source
+        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:
-        # Handle any other unexpected errors.
-        print(f"❌ Unexpected error connecting to {host}: {e}")
-        if transport:
-            transport.close()
-            
-def upload_to_ftp(filename):
-    """
-    Attempts to connect to two FTP hosts simultaneously and uses the first one that responds.
-    It now uses a corresponding directory for the primary and fallback hosts.
-    """
-    hosts = [FTP_HOST]
-    if FTP_HOST_FALLBACK:
-        hosts.append(FTP_HOST_FALLBACK)
-    
-    result_queue = queue.Queue()
-    threads = []
-    
-    print(f"--> Racing connections to {hosts} for uploading {filename}...")
-    for host in hosts:
-        thread = threading.Thread(target=connect_worker, args=(host, result_queue))
-        thread.daemon = True
-        thread.start()
-        threads.append(thread)
+        print(f"❌ Failed to initialize GCS client: {e}")
+
+# --- New GCS Upload Function (runs on CPU) ---
+def upload_to_gcs(image_object, filename):
+    """Uploads a PIL Image object to GCS from memory."""
+    if not gcs_client:
+        print("⚠️ GCS client not initialized. Skipping upload.")
+        return
+
     try:
-        winning_transport = result_queue.get(timeout=7)
+        print(f"--> Starting GCS upload for {filename}...")
+        bucket = gcs_client.bucket(GCS_BUCKET_NAME)
+        blob = bucket.blob(f"stablediff/{filename}")
+
+        # Convert PIL image to bytes stream
+        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()
         
-        # --- THIS IS THE NEW LOGIC ---
-        # 1. Determine which host won the race.
-        winning_host = winning_transport.getpeername()[0]
-        
-        # 2. Select the correct destination directory based on the winning host.
-        #    If the fallback directory isn't specified, it safely defaults to the primary directory.
-        if winning_host == FTP_HOST:
-            destination_directory = FTP_DIR
-        else:
-            destination_directory = FTP_DIR_FALLBACK if FTP_DIR_FALLBACK else FTP_DIR
-        print(f"--> Proceeding with upload to {winning_host} in directory {destination_directory}...")
-        # 3. Construct the full destination path using the selected directory.
-        sftp = paramiko.SFTPClient.from_transport(winning_transport)
-        destination_path = os.path.join(destination_directory, os.path.basename(filename))
-        sftp.put(filename, destination_path)
-        
-        print(f"✅ Successfully uploaded {filename}.")
-        
-        sftp.close()
-        winning_transport.close()
-    except queue.Empty:
-        print("❌ Critical Error: Neither FTP host responded in time.")
+        # Upload from the in-memory string
+        blob.upload_from_string(img_byte_arr, content_type='image/png')
+        print(f"✅ Successfully uploaded {filename} to GCS.")
+
     except Exception as e:
-        print(f"❌ An unexpected error occurred during SFTP operation: {e}")
+        print(f"❌ An error occurred during GCS upload: {e}")
 
+# --- Model Loading ---
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+hftoken = os.getenv("HF_AUTH_TOKEN")
 
-#vae=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", use_safetensors=True, subfolder='vae',token=True)
-vaeX=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", safety_checker=None, use_safetensors=True, subfolder='vae', low_cpu_mem_usage=False, torch_dtype=torch.float32)
 pipe = StableDiffusion3Pipeline.from_pretrained(
-    #"stabilityai  #  stable-diffusion-3.5-large",
     "ford442/stable-diffusion-3.5-large-bf16",
     trust_remote_code=True,
-    #vae=None,
-    #vae=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", use_safetensors=True, subfolder='vae',token=True),
-     #scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained('ford442/stable-diffusion-3.5-large-bf16', subfolder='scheduler',token=True),
-    #text_encoder=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
-    # text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
-    #text_encoder_2=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
-  #  text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
-    #text_encoder_3=None, #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
-  #  text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
-    #tokenizer=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer", token=True),
-    #tokenizer_2=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer_2", token=True),
-    transformer=None,
-    #tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=False, use_fast=True, subfolder="tokenizer_3", token=True),
-    #torch_dtype=torch.bfloat16,
+    transformer=None, # Load transformer separately
     use_safetensors=True,
+  #  token=hftoken
 )
-#text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
-#text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
-#text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
-ll_transformer=SD3Transformer2DModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='transformer').to(torch.device("cuda:0"), dtype=torch.bfloat16)
+# Load transformer separately and move to device with specified dtype
+ll_transformer=SD3Transformer2DModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='transformer', token=hftoken).to(device, dtype=torch.bfloat16)
 pipe.transformer=ll_transformer
 pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")
-
-#pipe.to(accelerator.device)
 pipe.to(device=device, dtype=torch.bfloat16)
 
-upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cuda'))
+upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-
 MAX_IMAGE_SIZE = 4096
 
-@spaces.GPU(duration=70)
-def infer_60(
-    prompt,
-    negative_prompt_1,
-    negative_prompt_2,
-    negative_prompt_3,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device='cuda').manual_seed(seed)
-    print('-- generating image --')
-    sd_image = pipe(
-            prompt=prompt,
-            prompt_2=prompt,
-            prompt_3=prompt,
-            negative_prompt=negative_prompt_1,
-            negative_prompt_2=negative_prompt_2,
-            negative_prompt_3=negative_prompt_3,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            max_sequence_length=512
-    ).images[0]
-    print('-- got image --')
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    sd35_path = f"sd35ll_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    upload_to_ftp(sd35_path)
-    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 --')
-    downscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale.save(upscale_path,optimize=False,compress_level=0)
-    upload_to_ftp(upscale_path)
-    return sd_image, prompt
-
-@spaces.GPU(duration=100)
-def infer_90(
-    prompt,
-    negative_prompt_1,
-    negative_prompt_2,
-    negative_prompt_3,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device='cuda').manual_seed(seed)
-    print('-- generating image --')
-    sd_image = pipe(
-            prompt=prompt,
-            prompt_2=prompt,
-            prompt_3=prompt,
-            negative_prompt=negative_prompt_1,
-            negative_prompt_2=negative_prompt_2,
-            negative_prompt_3=negative_prompt_3,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            max_sequence_length=512
-    ).images[0]
-    print('-- got image --')
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    sd35_path = f"sd35ll_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    upload_to_ftp(sd35_path)
-    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 --')
-    downscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale.save(upscale_path,optimize=False,compress_level=0)
-    upload_to_ftp(upscale_path)
-    return sd_image, prompt
-
+# --- Refactored GPU Inference Function ---
 @spaces.GPU(duration=120)
-def infer_110(
-    prompt,
-    negative_prompt_1,
-    negative_prompt_2,
-    negative_prompt_3,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
+def generate_images(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress=gr.Progress(track_tqdm=True)):
+    """Generates the main image and its upscaled version on the GPU."""
     seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device='cuda').manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
     print('-- generating image --')
     sd_image = pipe(
-            prompt=prompt,
-            prompt_2=prompt,
-            prompt_3=prompt,
-            negative_prompt=negative_prompt_1,
-            negative_prompt_2=negative_prompt_2,
-            negative_prompt_3=negative_prompt_3,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            max_sequence_length=512
+        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=512
     ).images[0]
     print('-- got image --')
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    sd35_path = f"sd35ll_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    upload_to_ftp(sd35_path)
+    
     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 --')
-    downscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale.save(upscale_path,optimize=False,compress_level=0)
-    upload_to_ftp(upscale_path)
-    return sd_image, prompt
+    downscaled_upscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4), Image.LANCZOS)
     
+    return sd_image, downscaled_upscale, prompt
+
+# --- Main Gradio Handler (runs on CPU) ---
+def run_inference_and_upload(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, save_consent, progress=gr.Progress(track_tqdm=True)):
+    """
+    Orchestrates the process: calls the GPU function, then handles the upload if consented.
+    """
+    # 1. Call the GPU-bound function to get the images
+    sd_image, upscaled_image, expanded_prompt = generate_images(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress)
+
+    # 2. If user consented, start uploads in background threads
+    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"
+
+        # Create and start threads for each upload
+        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.")
+
+    # 3. Return the primary image to the UI immediately
+    return sd_image, expanded_prompt
+
+# --- Gradio UI Definition ---
 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.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)  # Add this line
+        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,
+                label="Prompt", show_label=False, max_lines=1,
+                placeholder="Enter your prompt", container=False,
             )
-            run_button_60 = gr.Button("Run 60", scale=0, variant="primary")
-            run_button_90 = gr.Button("Run 90", scale=0, variant="primary")
-            run_button_110 = gr.Button("Run 110", scale=0, variant="primary")
-        result = gr.Image(label="Result", show_label=False)
+            # Use a single run button for simplicity or keep multiple if durations are critical
+            run_button = gr.Button("Run", scale=0, variant="primary")
+            
+        result = gr.Image(label="Result", show_label=False, type="pil")
+        
+        # --- New Consent Checkbox ---
+        save_consent_checkbox = gr.Checkbox(
+            label="✅ Anonymously upload result to a public gallery",
+            value=False, # 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",
-                visible=True,
-                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",
-                visible=True,
-                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",
-                visible=True,
-                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)"
-            )
-            num_iterations = gr.Number(
-                value=1000, 
-                label="Number of Iterations")
+            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=768,
-                )
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=768,
-                )
-                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="Number of inference steps",
-                    minimum=1,
-                    maximum=500,
-                    step=1,
-                    value=100,
-                )
-        gr.on(
-        triggers=[run_button_60.click, prompt.submit],
-        fn=infer_60,
-        inputs=[
-            prompt,
-            negative_prompt_1,
-            negative_prompt_2,
-            negative_prompt_3,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, expanded_prompt_output],
-        )
-        gr.on(
-        triggers=[run_button_90.click, prompt.submit],
-        fn=infer_90,
-        inputs=[
-            prompt,
-            negative_prompt_1,
-            negative_prompt_2,
-            negative_prompt_3,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, expanded_prompt_output],
-        )
-        gr.on(
-        triggers=[run_button_110.click, prompt.submit],
-        fn=infer_110,
-        inputs=[
-            prompt,
-            negative_prompt_1,
-            negative_prompt_2,
-            negative_prompt_3,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, expanded_prompt_output],
+                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
+                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
+            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)
+
+        # Connect the button to the main handler function
+        run_button.click(
+            fn=run_inference_and_upload,
+            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__":