Update app.py

app.py

@@ -6,182 +6,96 @@ import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=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: 
-    # Load VAE component
-    vae = AutoencoderKL.from_pretrained(
-        "madebyollin/sdxl-vae-fp16-fix", 
-    )
+# Load VAE component
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", 
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
 
-    # Configure the pipeline
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "dataautogpt3/ProteusV0.5", 
-        vae=vae,
-    )
-    pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    pipe.to(device)
+# Configure the pipeline
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "dataautogpt3/ProteusV0.5", 
+    vae=vae,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
+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, clip_skip):
-
+def infer(prompt, negative_prompt="", seed=0, randomize_seed=True, width=1024, height=1024, guidance_scale=7.0, num_inference_steps=50, clip_skip=2):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
     
     image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator,
-        clip_skip = clip_skip
+        prompt=prompt, 
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale, 
+        num_inference_steps=num_inference_steps, 
+        width=width, 
+        height=height,
+        generator=generator,
+        clip_skip=clip_skip
     ).images[0] 
     
     return image
 
 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"
+    ["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="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
-}
+css = """
+#col-container { max-width: 700px; margin: 0 auto; }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
+demo = gr.Blocks(css=css)
 
-with gr.Blocks(css=css) as demo:
+with demo:
+    gr.Markdown(f"# ProteusV0.5 Demo")
+    gr.Markdown(f"Currently running on {'GPU' if torch.cuda.is_available() else 'CPU'}.")
+    
+    gr.Markdown("""
+    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.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # 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():
-            
-            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.Image(label="Result", show_label=False)
+        prompt = gr.Text(label="Prompt", placeholder="Enter your prompt")
+        run_button = gr.Button("Generate Image")
+        result = gr.Image(label="Generated Image")
 
         with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
+            negative_prompt = gr.Text(label="Negative prompt", placeholder="Enter a negative prompt")
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=512,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=512,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=7.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=50,
-                )
+            width = gr.Slider(label="Width", minimum=512, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+            height = gr.Slider(label="Height", minimum=512, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+            guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=20.0, step=0.1, value=7.0)
+            num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=100, step=1, value=50)
+            clip_skip = gr.Slider(label="Clip skip", minimum=1, maximum=12, step=1, value=2)
 
-            clip_skip = gr.Slider(
-                label="Clip skip",
-                minimum=1,
-                maximum=12,
-                step=1,
-                value=2,
-            )
-        
         gr.Examples(
             examples=examples,
-            inputs=[prompt],
-            outputs=[result],
+            inputs=prompt,
+            outputs=result,
             fn=infer,
             cache_examples=True
-    )
+        )
 
     run_button.click(
         fn=infer,
         inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, clip_skip],
-        outputs=[result]
+        outputs=result
     )
 
 demo.queue().launch()