Update app.py

app.py

@@ -12,7 +12,7 @@ import spaces
 import torch
 from diffusers import AutoencoderKL, DiffusionPipeline
 
-DESCRIPTION = "# SDXL"
+DESCRIPTION = "humblemikey/PixelWave10"
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
@@ -20,11 +20,10 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
-ENABLE_REFINER = os.getenv("ENABLE_REFINER", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+    #vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
     pipe = DiffusionPipeline.from_pretrained(
         "humblemikey/PixelWave10",
         #vae=vae,
@@ -32,14 +31,6 @@ if torch.cuda.is_available():
         use_safetensors=True,
         #variant="fp16",
     )
-    if ENABLE_REFINER:
-        refiner = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-refiner-1.0",
-            vae=vae,
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            variant="fp16",
-        )
 
     if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
@@ -47,13 +38,9 @@ if torch.cuda.is_available():
             refiner.enable_model_cpu_offload()
     else:
         pipe.to(device)
-        if ENABLE_REFINER:
-            refiner.to(device)
 
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        if ENABLE_REFINER:
-            refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
 
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
@@ -66,67 +53,31 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 def generate(
     prompt: str,
     negative_prompt: str = "",
-    prompt_2: str = "",
-    negative_prompt_2: str = "",
     use_negative_prompt: bool = False,
-    use_prompt_2: bool = False,
-    use_negative_prompt_2: bool = False,
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
-    guidance_scale_base: float = 5.0,
-    guidance_scale_refiner: float = 5.0,
-    num_inference_steps_base: int = 25,
-    num_inference_steps_refiner: int = 25,
-    apply_refiner: bool = False,
+    guidance_scale_base: float = 4.0,
+    num_inference_steps_base: int = 40,
 ) -> PIL.Image.Image:
     generator = torch.Generator().manual_seed(seed)
 
     if not use_negative_prompt:
         negative_prompt = None  # type: ignore
-    if not use_prompt_2:
-        prompt_2 = None  # type: ignore
-    if not use_negative_prompt_2:
-        negative_prompt_2 = None  # type: ignore
-
-    if not apply_refiner:
-        return pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale_base,
-            num_inference_steps=num_inference_steps_base,
-            generator=generator,
-            output_type="pil",
-        ).images[0]
-    else:
-        latents = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale_base,
-            num_inference_steps=num_inference_steps_base,
-            generator=generator,
-            output_type="latent",
-        ).images
-        image = refiner(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            guidance_scale=guidance_scale_refiner,
-            num_inference_steps=num_inference_steps_refiner,
-            image=latents,
-            generator=generator,
-        ).images[0]
-        return image
 
+    return pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        prompt_2=prompt_2,
+        negative_prompt_2=negative_prompt_2,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale_base,
+        num_inference_steps=num_inference_steps_base,
+        generator=generator,
+        output_type="pil",
+    ).images[0]
+    return image
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
@@ -154,27 +105,12 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
             use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
-            use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
-            use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False)
         negative_prompt = gr.Text(
             label="Negative prompt",
             max_lines=1,
             placeholder="Enter a negative prompt",
             visible=False,
         )
-        prompt_2 = gr.Text(
-            label="Prompt 2",
-            max_lines=1,
-            placeholder="Enter your prompt",
-            visible=False,
-        )
-        negative_prompt_2 = gr.Text(
-            label="Negative prompt 2",
-            max_lines=1,
-            placeholder="Enter a negative prompt",
-            visible=False,
-        )
-
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -198,7 +134,6 @@ with gr.Blocks(css="style.css") as demo:
                 step=32,
                 value=1024,
             )
-        apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
         with gr.Row():
             guidance_scale_base = gr.Slider(
                 label="Guidance scale for base",
@@ -214,21 +149,6 @@ with gr.Blocks(css="style.css") as demo:
                 step=1,
                 value=25,
             )
-        with gr.Row(visible=False) as refiner_params:
-            guidance_scale_refiner = gr.Slider(
-                label="Guidance scale for refiner",
-                minimum=1,
-                maximum=20,
-                step=0.1,
-                value=5.0,
-            )
-            num_inference_steps_refiner = gr.Slider(
-                label="Number of inference steps for refiner",
-                minimum=10,
-                maximum=100,
-                step=1,
-                value=25,
-            )
 
     gr.Examples(
         examples=examples,
@@ -244,34 +164,11 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     )
-    use_prompt_2.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_prompt_2,
-        outputs=prompt_2,
-        queue=False,
-        api_name=False,
-    )
-    use_negative_prompt_2.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt_2,
-        outputs=negative_prompt_2,
-        queue=False,
-        api_name=False,
-    )
-    apply_refiner.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=apply_refiner,
-        outputs=refiner_params,
-        queue=False,
-        api_name=False,
-    )
 
     gr.on(
         triggers=[
             prompt.submit,
             negative_prompt.submit,
-            prompt_2.submit,
-            negative_prompt_2.submit,
             run_button.click,
         ],
         fn=randomize_seed_fn,
@@ -284,19 +181,13 @@ with gr.Blocks(css="style.css") as demo:
         inputs=[
             prompt,
             negative_prompt,
-            prompt_2,
-            negative_prompt_2,
             use_negative_prompt,
-            use_prompt_2,
-            use_negative_prompt_2,
             seed,
             width,
             height,
             guidance_scale_base,
             guidance_scale_refiner,
             num_inference_steps_base,
-            num_inference_steps_refiner,
-            apply_refiner,
         ],
         outputs=result,
         api_name="run",