Update app.py

app.py

@@ -1,22 +1,32 @@
-#!/usr/bin/env python
 
-import os
-import random
-import uuid
+if torch.cuda.is_available():
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "RunDiffusion/Juggernaut-X-v10",
+        torch_dtype=torch.float16
+    )
+    #pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    
+    
+    #pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
+    #pipe.set_adapters("dalle")
+
+    pipe.to("cuda")
+    
+    
 
 import gradio as gr
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler, LCMScheduler, AutoencoderKL,DiffusionPipeline
+import torch
 import numpy as np
-from PIL import Image
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 import spaces
-import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-
-DESCRIPTION = """
-# Juggernaut X v10
-"""
+import os
+import random
+import uuid
 
 def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".jpeg"
+    unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
@@ -26,34 +36,37 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 MAX_SEED = np.iinfo(np.int32).max
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
-
-MAX_SEED = np.iinfo(np.int32).max
 
-USE_TORCH_COMPILE = 0
-ENABLE_CPU_OFFLOAD = 0
+JX_pipe = StableDiffusionXLPipeline.from_pretrained(
+        "RunDiffusion/Juggernaut-X-Hyper",
+        vae=vae,
+        torch_dtype=torch.float16,
+    )
+JX_pipe.to("cuda")
 
 
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
+J10_pipe = StableDiffusionXLPipeline.from_pretrained(
         "RunDiffusion/Juggernaut-X-v10",
-        torch_dtype=torch.float16
+        vae=vae,
+        torch_dtype=torch.float16,
     )
-    #pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    
-    
-    #pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
-    #pipe.set_adapters("dalle")
+J10_pipe.to("cuda")
 
-    pipe.to("cuda")
-    
-    
 
-@spaces.GPU(enable_queue=True)
-def generate(
-    prompt: str,
+J9_pipe = StableDiffusionXLPipeline.from_pretrained(
+        "RunDiffusion/Juggernaut-X-v10",
+        vae=vae,
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+        add_watermarker=False,
+        variant="fp16"
+    )
+J9_pipe.to("cuda")
+
+@spaces.GPU
+def run_comparison(prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
     num_inference_steps: int = 30,
@@ -65,15 +78,22 @@ def generate(
     randomize_seed: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ):
-
-    
     seed = int(randomize_seed_fn(seed, randomize_seed))
-
     if not use_negative_prompt:
-        negative_prompt = ""  # type: ignore
+        negative_prompt = ""
+    image_r3 = JX_pipe(prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=num_images_per_prompt,
+        cross_attention_kwargs={"scale": 0.65},
+        output_type="pil",
+    ).images
+    image_paths_r3 = [save_image(img) for img in image_r3]
 
-    images = pipe(
-        prompt=prompt,
+    image_r4 = JX10_pipe(prompt=prompt,
         negative_prompt=negative_prompt,
         width=width,
         height=height,
@@ -83,42 +103,35 @@ def generate(
         cross_attention_kwargs={"scale": 0.65},
         output_type="pil",
     ).images
-    image_paths = [save_image(img) for img in images]
-    print(image_paths)
-    return image_paths, seed
-
-examples = [
-    "neon holography crystal cat",
-    "a cat eating a piece of cheese",
-    "an astronaut riding a horse in space",
-    "a cartoon of a boy playing with a tiger",
-    "a cute robot artist painting on an easel, concept art",
-    "a close up of a woman wearing a transparent, prismatic, elaborate nemeses headdress, over the should pose, brown skin-tone"
-]
+    image_paths_r4 = [save_image(img) for img in image_r4]
 
-css = '''
-.gradio-container{max-width: 560px !important}
-h1{text-align:center}
-'''
-with gr.Blocks(css=css) as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=False,
-    )
+    image_r5 = JX9_pipe(prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=num_images_per_prompt,
+        cross_attention_kwargs={"scale": 0.65},
+        output_type="pil",
+    ).images
+    image_paths_r5 = [save_image(img) for img in image_r5]
+    return image_paths_r3, image_paths_r4,image_paths_r5, seed
+
+examples = ["A dignified beaver wearing glasses, a vest, and colorful neck tie.",
+"The spirit of a tamagotchi wandering in the city of Barcelona",
+"an ornate, high-backed mahogany chair with a red cushion",
+"a sketch of a camel next to a stream",
+"a delicate porcelain teacup sits on a saucer, its surface adorned with intricate blue patterns",
+"a baby swan grafitti",
+"A bald eagle made of chocolate powder, mango, and whipped cream"
+]
 
-    with gr.Group():
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Run", scale=0)
-        result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
+with gr.Blocks() as demo:
+    gr.Markdown("## One step SDXL comparison 🦶")
+    gr.Markdown('Compare SDXL variants and distillations able to generate images in a single diffusion step')
+    prompt = gr.Textbox(label="Prompt")
+    run = gr.Button("Run")
     with gr.Accordion("Advanced options", open=False):
         use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
         negative_prompt = gr.Text(
@@ -178,46 +191,47 @@ with gr.Blocks(css=css) as demo:
                 value=6,
             )
 
-        
-
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=[result, seed],
-        fn=generate,
-        cache_examples=False,
+    with gr.Row():
+        with gr.Column():
+            image_r3 = gr.Gallery(label="RealVisXL V3",columns=1, preview=True,)
+            gr.Markdown("## [RealVisXL V3](https://huggingface.co)")
+        with gr.Column():
+            image_r4 = gr.Gallery(label="RealVisXL V4",columns=1, preview=True,)
+            gr.Markdown("## [RealVisXL V4](https://huggingface.co)")
+        with gr.Column():
+            image_r5 = gr.Gallery(label="Playground v2.5",columns=1, preview=True,)
+            gr.Markdown("## [Playground v2.5](https://huggingface.co)")    
+    image_outputs = [image_r3, image_r4, image_r5]
+    gr.on(
+        triggers=[prompt.submit, run.click],
+        fn=run_comparison,
+        inputs=[
+           prompt,
+           negative_prompt,
+           use_negative_prompt,
+           num_inference_steps,
+           num_images_per_prompt,
+           seed,
+           width,
+           height,
+           guidance_scale,
+           randomize_seed,
+       ],
+        outputs=image_outputs
     )
-
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
         outputs=negative_prompt,
         api_name=False,
     )
-    
-
-    gr.on(
-        triggers=[
-            prompt.submit,
-            negative_prompt.submit,
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            num_inference_steps,
-            num_images_per_prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            randomize_seed,
-        ],
-        outputs=[result, seed],
-        api_name="run",
+    gr.Examples(
+        examples=examples,
+        fn=run_comparison,
+        inputs=prompt,
+        outputs=image_outputs,
+        cache_examples=False,
+        run_on_click=True
     )
-    
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(show_api=False, debug=False)