Add 4x upscaling with stabilityai/stable-diffusion-x4-upscaler

- Load StableDiffusionUpscalePipeline for 4x upscaling
- Add enable_upscale checkbox in Advanced Settings
- Apply upscaling as final step after image generation
- Maintain live preview during generation, then upscale final image
- Add error handling for upscaling failures

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import random
 import spaces
 import torch
-from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL, StableDiffusionUpscalePipeline
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from huggingface_hub import hf_hub_download
@@ -17,6 +17,9 @@ taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
 
+# Load upscaler pipeline
+upscaler = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=dtype).to(device)
+
 # Available LoRAs
 LORAS = {
     "None": None,
@@ -69,7 +72,7 @@ MAX_IMAGE_SIZE = 2048
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 @spaces.GPU(duration=75)
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, lora_selection="None", progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, lora_selection="None", enable_upscale=False, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
@@ -86,6 +89,7 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
         print(f"Failed to load LoRA {lora_selection}: {e}")
     
     try:
+        final_img = None
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
                 prompt=prompt,
                 guidance_scale=guidance_scale,
@@ -96,7 +100,24 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
                 output_type="pil",
                 good_vae=good_vae,
             ):
+                final_img = img
                 yield img, seed
+        
+        # Apply upscaling if enabled
+        if enable_upscale and final_img is not None:
+            try:
+                upscaled_img = upscaler(
+                    prompt=prompt,
+                    image=final_img,
+                    num_inference_steps=20,
+                    guidance_scale=7.5,
+                    generator=generator,
+                ).images[0]
+                yield upscaled_img, seed
+            except Exception as e:
+                print(f"Error during upscaling: {e}")
+                yield final_img, seed
+                
     except Exception as e:
         print(f"Error during generation: {e}")
         # Fallback to basic generation
@@ -108,6 +129,20 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
             height=height,
             generator=generator,
         ).images[0]
+        
+        # Apply upscaling if enabled
+        if enable_upscale:
+            try:
+                img = upscaler(
+                    prompt=prompt,
+                    image=img,
+                    num_inference_steps=20,
+                    guidance_scale=7.5,
+                    generator=generator,
+                ).images[0]
+            except Exception as e:
+                print(f"Error during upscaling: {e}")
+                
         yield img, seed
     
 examples = [
@@ -154,6 +189,12 @@ with gr.Blocks(css=css) as demo:
                 info="Select a LoRA to enhance image generation"
             )
             
+            enable_upscale = gr.Checkbox(
+                label="Enable 4x Upscaling",
+                value=False,
+                info="Upscale final image using Stable Diffusion 4x upscaler"
+            )
+            
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -211,7 +252,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_selection],
+        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_selection, enable_upscale],
         outputs = [result, seed]
     )