Update app.py

app.py

@@ -1,40 +1,35 @@
-import spaces
-import os
+import subprocess
+subprocess.run(['sh', './spaces.sh'])
 
- # 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["SAFETENSORS_FAST_GPU"] = "1"
-# os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
+import os
+# Environment variable setup
+os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
+os.environ['TORCH_LINALG_PREFER_CUSOLVER'] = '1'
+os.environ['PYTORCH_CUDA_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 spaces
 import gradio as gr
 import numpy as np
 import random
+import datetime
+import threading
+import io
 
-import torch
-from diffusers import StableDiffusion3Pipeline
-from transformers import CLIPTextModelWithProjection, T5EncoderModel
-from transformers import CLIPTokenizer, T5TokenizerFast
+# --- New GCS Imports ---
+from google.oauth2 import service_account
+from google.cloud import storage
 
-import re
-import paramiko
-import urllib
-import time
-from image_gen_aux import UpscaleWithModel
-from huggingface_hub import hf_hub_download
-import datetime
-#import cyper
+import torch
 
-from diffusers import AutoencoderKL
-#from models.transformer_sd3 import SD3Transformer2DModel
-#from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
+@spaces.GPU(required=True)
+def install_flashattn():
+    subprocess.run(['sh', './flashattn.sh'])
 
-from PIL import Image
+#install_flashattn()
 
+# 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
@@ -45,424 +40,262 @@ torch.backends.cuda.preferred_blas_library="cublas"
 torch.backends.cuda.preferred_linalg_library="cusolver"
 torch.set_float32_matmul_precision("highest")
 
-hftoken = os.getenv("HF_AUTH_TOKEN") 
-
-# code = r'''
-import torch
-import paramiko
-import os
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, AutoencoderKL
+from PIL import Image
+from image_gen_aux import UpscaleWithModel
 
-FTP_HOST = 'noahcohn.com'
-FTP_USER = 'ford442'
-FTP_PASS = os.getenv("FTP_PASS")
-FTP_DIR = 'img.noahcohn.com/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
 
-def upload_to_ftp(filename):
+# Initialize GCS client if credentials are available
+gcs_client = None
+if GCS_SA_KEY and GCS_BUCKET_NAME:
     try:
-        transport = paramiko.Transport((FTP_HOST, 22))
-        destination_path=FTP_DIR+filename
-        transport.connect(username = FTP_USER, password = FTP_PASS)
-        sftp = paramiko.SFTPClient.from_transport(transport)
-        sftp.put(filename, destination_path)
-        sftp.close()
-        transport.close()
-        print(f"Uploaded {filename} to FTP server")
+        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:
-        print(f"FTP upload error: {e}")
-# '''
+        print(f"❌ Failed to initialize GCS client: {e}")
 
-# pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))
+def upload_to_gcs(image_object, filename):
+    if not gcs_client:
+        print("⚠️ GCS client not initialized. Skipping upload.")
+        return
+    try:
+        print(f"--> Starting GCS upload for {filename}...")
+        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.")
+    except Exception as e:
+        print(f"❌ An error occurred during GCS upload: {e}")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    #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, token=True)
- 
-pipe = StableDiffusion3Pipeline.from_single_file(
-    "https://huggingface.co/1inkus/sd35-large-UltraReal-bf16-DDUF/blob/main/sd3-bf16-large.dduf",
-    #vae=None,
-    #tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=False, use_fast=True, subfolder="tokenizer_3"),
-    use_safetensors=True,
-    config = 'ford442/stable-diffusion-3.5-large-bf16',
-    #dtype=torch.bfloat16,
-    #devive_map='auto',
-) #.to(device=device)
 
-###    pipe = StableDiffusion3Pipeline.from_pretrained(
-    #"stabilityai  #  stable-diffusion-3.5-large",
-    #"ford442/stable-diffusion-3.5-large-bf16",
-###    "ford442/stable-diffusion-3.5-large-fp32",
-    #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),
-    #tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=False, use_fast=True, subfolder="tokenizer_3", token=True),
-###    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,
-    #use_safetensors=False,
-###)
-#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-fp32", subfolder='text_encoder_3',token=True).to(torch.device("cuda:0")) #, dtype=torch.bfloat16)
-    
-#pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")
+def load_model():
+    pipe = StableDiffusion3Pipeline.from_pretrained(
+        "ford442/stable-diffusion-3.5-large-bf16",
+        trust_remote_code=True,
+        transformer=None, # Load transformer separately
+        use_safetensors=True
+    )
+    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)
+    upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(device)
+    return pipe, upscaler_2
 
-#pipe.to(device=device) #, dtype=torch.bfloat16)
-#pipe.vae=vaeX #.to('cpu')
-
-pipe.to(device,torch.bfloat16)
-upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cuda'))
+pipe, upscaler_2 = load_model()
 
 MAX_SEED = np.iinfo(np.int32).max
-
 MAX_IMAGE_SIZE = 4096
 
-@spaces.GPU(duration=40)
-def infer_30(
-    prompt,
-    negative_prompt_1,
-    negative_prompt_2,
-    negative_prompt_3,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    #pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
+
+@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='cuda').manual_seed(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=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,
-         #   cross_attention_kwargs={"scale": 0.75},
-            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=384
     ).images[0]
     print('-- got image --')
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    sd35_path = f"sd35ll32_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(sd35_path)
-    upload_to_ftp(sd35_path)
-    #  pipe.unet.to('cpu')
-    #upscaler_2.to(torch.device('cuda'))
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     with torch.no_grad():
-        upscale2 = upscaler_2(sd_image, 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 --')
-    downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale2.save(upscale_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(upscale_path)
-    upload_to_ftp(upscale_path)
-    return sd_image, prompt
+    downscaled_upscale = upscale2.resize((upscale2.width // 16, upscale2.height // 16), Image.LANCZOS)
+    return sd_image, downscaled_upscale, prompt
 
 @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),
-):
-    #pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
+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='cuda').manual_seed(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=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,
-         #   cross_attention_kwargs={"scale": 0.75},
-            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=384
     ).images[0]
     print('-- got image --')
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    sd35_path = f"sd35ll32_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(sd35_path)
-    upload_to_ftp(sd35_path)
-    #  pipe.unet.to('cpu')
-    #upscaler_2.to(torch.device('cuda'))
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     with torch.no_grad():
-        upscale2 = upscaler_2(sd_image, 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 --')
-    downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale2.save(upscale_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(upscale_path)
-    upload_to_ftp(upscale_path)
-    return sd_image, prompt
+    downscaled_upscale = upscale2.resize((upscale2.width // 16, upscale2.height // 16), Image.LANCZOS)
+    return sd_image, downscaled_upscale, 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),
-):
-    #pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
+@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='cuda').manual_seed(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=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,
-         #   cross_attention_kwargs={"scale": 0.75},
-            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=384
     ).images[0]
     print('-- got image --')
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    sd35_path = f"sd35ll32_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(sd35_path)
-    upload_to_ftp(sd35_path)
-    #  pipe.unet.to('cpu')
-    #upscaler_2.to(torch.device('cuda'))
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     with torch.no_grad():
-        upscale2 = upscaler_2(sd_image, 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 --')
-    downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale2.save(upscale_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(upscale_path)
-    upload_to_ftp(upscale_path)
-    return sd_image, prompt
+    downscaled_upscale = upscale2.resize((upscale2.width // 16, upscale2.height // 16), Image.LANCZOS)
+    return sd_image, downscaled_upscale, prompt
 
-@spaces.GPU(duration=110)
-def infer_100(
-    prompt,
-    negative_prompt_1,
-    negative_prompt_2,
-    negative_prompt_3,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    #pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
-    #pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
-    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,
-         #   cross_attention_kwargs={"scale": 0.75},
-            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"sd35ll32_{timestamp}.png"
-    sd_image.save(sd35_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(sd35_path)
-    upload_to_ftp(sd35_path)
-    #  pipe.unet.to('cpu')
-    #upscaler_2.to(torch.device('cuda'))
-    with torch.no_grad():
-        upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
-    print('-- got upscaled image --')
-    downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35ll_upscale_{timestamp}.png"
-    downscale2.save(upscale_path,optimize=False,compress_level=0)
-    #pyx.upload_to_ftp(upscale_path)
-    upload_to_ftp(upscale_path)
-    return sd_image, 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
+
+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;}
 """
-
-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")
-        expanded_prompt_output = gr.Textbox(label="Prompt", lines=1)  # Add this line
+        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,
+                label="Prompt", show_label=False, max_lines=1,
+                placeholder="Enter your prompt", container=False,
             )
-            run_button_30 = gr.Button("Run 30", scale=0, variant="primary")
-            run_button_60 = gr.Button("Run 60", scale=0, variant="primary")
-            run_button_90 = gr.Button("Run 90", scale=0, variant="primary")
-            run_button_100 = gr.Button("Run 100", scale=0, variant="primary")
-        result = gr.Image(label="Result", show_label=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",
-                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=50,
-                )
-        gr.on(
-        triggers=[run_button_30.click, prompt.submit],
-        fn=infer_30,
-        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_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],
+                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],
         )
-        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],
+
+        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],
         )
-        gr.on(
-        triggers=[run_button_100.click, prompt.submit],
-        fn=infer_100,
-        inputs=[
-            prompt,
-            negative_prompt_1,
-            negative_prompt_2,
-            negative_prompt_3,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        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()