Update app.py

app.py

@@ -51,43 +51,32 @@ if GCS_SA_KEY and GCS_BUCKET_NAME:
     except Exception as e:
         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:
         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()
-        
-        # 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 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")
 
 pipe = StableDiffusion3Pipeline.from_pretrained(
     "ford442/stable-diffusion-3.5-large-bf16",
     trust_remote_code=True,
     transformer=None, # Load transformer separately
-    use_safetensors=True,
-  #  token=hftoken
+    use_safetensors=True
 )
-# 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)
+
+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)
@@ -97,13 +86,10 @@ upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(devic
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
-# --- Refactored GPU Inference Function ---
-@spaces.GPU(duration=120)
-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."""
+@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 --')
     sd_image = pipe(
         prompt=prompt, prompt_2=prompt, prompt_3=prompt,
@@ -113,48 +99,102 @@ def generate_images(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, hei
         max_sequence_length=512
     ).images[0]
     print('-- got image --')
-    
     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 // 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)
+@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 --')
+    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=512
+    ).images[0]
+    print('-- got image --')
+    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 // 4, upscale2.height // 4), Image.LANCZOS)
+    return sd_image, downscaled_upscale, prompt
 
-    # 2. If user consented, start uploads in background threads
+@spaces.GPU(duration=110)
+def generate_images_100(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 --')
+    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=512
+    ).images[0]
+    print('-- got image --')
+    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 // 4, upscale2.height // 4), Image.LANCZOS)
+    return sd_image, downscaled_upscale, prompt
+
+def run_inference_and_upload_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, save_consent, progress=gr.Progress(track_tqdm=True)):
+    sd_image, upscaled_image, expanded_prompt = generate_images_30(prompt, neg_prompt_1, neg_prompt_2, neg_prompt_3, width, height, guidance, steps, progress)
     if save_consent:
         print("✅ User consented to save. Preparing uploads...")
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         sd_filename = f"sd35ll_{timestamp}.png"
         upscale_filename = f"sd35ll_upscale_{timestamp}.png"
-
-        # 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.")
+    return sd_image, expanded_prompt
 
-    # 3. Return the primary image to the UI immediately
+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
 
-# --- Gradio UI Definition ---
+def run_inference_and_upload_100(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_100(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.Column(elem_id="col-container"):
         gr.Markdown(" # StableDiffusion 3.5 Large with UltraReal lora test")
@@ -164,18 +204,15 @@ with gr.Blocks(theme=gr.themes.Origin(), css=css) as demo:
                 label="Prompt", show_label=False, max_lines=1,
                 placeholder="Enter your prompt", container=False,
             )
-            # Use a single run button for simplicity or keep multiple if durations are critical
-            run_button = gr.Button("Run", scale=0, variant="primary")
-            
+            run_button_30 = gr.Button("Run30", scale=0, variant="primary")
+            run_button_60 = gr.Button("Run60", scale=0, variant="primary")
+            run_button_100 = gr.Button("Run100", 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
+            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)")
@@ -187,9 +224,8 @@ 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)
 
-        # Connect the button to the main handler function
-        run_button.click(
-            fn=run_inference_and_upload,
+        run_button_30.click(
+            fn=run_inference_and_upload_30,
             inputs=[
                 prompt,
                 negative_prompt_1,
@@ -204,5 +240,38 @@ with gr.Blocks(theme=gr.themes.Origin(), css=css) as demo:
             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_100.click(
+            fn=run_inference_and_upload_100,
+            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()