Update app.py

app.py

@@ -1,24 +1,45 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionXLPipeline, KDPM2AncestralDiscreteScheduler, AutoencoderKL
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 if torch.cuda.is_available():
     torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
+    # Load VAE component
+    vae = AutoencoderKL.from_pretrained(
+        "madebyollin/sdxl-vae-fp16-fix", 
+        torch_dtype=torch.float16
+    )
+
+    # Configure the pipeline
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "dataautogpt3/ProteusV0.5", 
+        vae=vae,
+        torch_dtype=torch.float16
+    )
+    pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe.to(device)
 else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+    # Load VAE component
+    vae = AutoencoderKL.from_pretrained(
+        "madebyollin/sdxl-vae-fp16-fix", 
+    )
+
+    # Configure the pipeline
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "dataautogpt3/ProteusV0.5", 
+        vae=vae,
+    )
+    pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, clip_skip):
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -32,7 +53,8 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         num_inference_steps = num_inference_steps, 
         width = width, 
         height = height,
-        generator = generator
+        generator = generator,
+        clip_skip = clip_skip
     ).images[0] 
     
     return image
@@ -41,6 +63,8 @@ examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
     "A delicious ceviche cheesecake slice",
+    "black fluffy gorgeous dangerous cat animal creature, large orange eyes, big fluffy ears, piercing gaze, full moon, dark ambiance, best quality, extremely detailed",
+    "high quality pixel art, a pixel art silhouette of an anime space-themed girl in a space-punk steampunk style, lying in her bed by the window of a spaceship, smoking, with a rustic feel. The image should embody epic portraiture and double exposure, featuring an isolated landscape visible through the window. The colors should primarily be dynamic and action-packed, with a strong use of negative space. The entire artwork should be in pixel art style, emphasizing the characters shape and set against a white background. Silhouette"
 ]
 
 css="""
@@ -59,8 +83,14 @@ with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # Text-to-Image Gradio Template
+        # ProteusV0.5 Demo
         Currently running on {power_device}.
+        
+        ProteusV0.5 is a state-of-the-art text-to-image model that leverages the power of Stable Diffusion to generate high-quality images from text prompts. This model has been fine-tuned on a large dataset of images and has been trained to understand a wide range of prompts and styles.
+        
+        With ProteusV0.5, you can generate images in a variety of styles, from realistic to abstract, and from simple to complex. The model is also capable of understanding and responding to complex prompts, making it a powerful tool for artists, designers, and anyone looking to generate high-quality images.
+        
+        To use this demo, simply enter a text prompt in the input field below, and the model will generate an image based on your prompt. You can also adjust the settings to control the quality and style of the generated image.
         """)
         
         with gr.Row():
@@ -100,18 +130,18 @@ with gr.Blocks(css=css) as demo:
                 
                 width = gr.Slider(
                     label="Width",
-                    minimum=256,
+                    minimum=512,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
                 
                 height = gr.Slider(
                     label="Height",
-                    minimum=256,
+                    minimum=512,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
             
             with gr.Row():
@@ -121,16 +151,24 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,
+                    value=7.0,
                 )
                 
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
-                    maximum=12,
+                    maximum=50,
                     step=1,
-                    value=2,
+                    value=50,
                 )
+
+            clip_skip = gr.Slider(
+                label="Clip skip",
+                minimum=1,
+                maximum=12,
+                step=1,
+                value=2,
+            )
         
         gr.Examples(
             examples = examples,
@@ -139,7 +177,7 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, clip_skip],
         outputs = [result]
     )