Update app.py

app.py

@@ -3,14 +3,12 @@ import random
 import uuid
 import gradio as gr
 import numpy as np
-import torch
-import concurrent.futures
 from PIL import Image
+import spaces
+import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 from typing import Tuple
 
-
-# Function to apply styles
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     styles = {
         "3840 x 2160": (
@@ -24,8 +22,6 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + negative
 
-
-# Function to load and prepare the model
 def load_and_prepare_model():
     model_id = "SG161222/RealVisXL_V5.0_Lightning"
     pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -37,25 +33,19 @@ def load_and_prepare_model():
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     return pipe
 
-
 model = load_and_prepare_model()
 
-
-# Function to randomize or apply a fixed seed
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, np.iinfo(np.int32).max)
     return seed
 
-
-# Function to save generated images
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
-
-# Main generation function with timeout
+@spaces.GPU(duration=60, enable_queue=True)
 def generate(
         prompt: str,
         seed: int = 1,
@@ -67,37 +57,26 @@ def generate(
 ):
     global model
 
-    def run_model():
-        seed = int(randomize_seed_fn(seed, randomize_seed))
-        generator = torch.Generator(device=model.device).manual_seed(seed)
-
-        positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
-
-        options = {
-            "prompt": [positive_prompt],
-            "negative_prompt": [negative_prompt],
-            "width": width,
-            "height": height,
-            "guidance_scale": guidance_scale,
-            "num_inference_steps": num_inference_steps,
-            "generator": generator,
-            "output_type": "pil",
-        }
-
-        images = model(**options).images
-        return save_image(images[0])  # Save and return the image path
-
-    # Use a timeout for CPU-bound operations
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        future = executor.submit(run_model)
-        try:
-            image_path = future.result(timeout=180)  # Set timeout to 180 seconds
-            return image_path
-        except concurrent.futures.TimeoutError:
-            return "Error: Image generation timed out. Please try again with different settings or on a GPU-enabled setup."
-
-
-# Gradio interface setup
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator(device=model.device).manual_seed(seed)
+
+    positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
+
+    options = {
+        "prompt": [positive_prompt],
+        "negative_prompt": [negative_prompt],
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "pil",
+    }
+
+    images = model(**options).images
+    image_path = save_image(images[0])  # Saving the first generated image
+    return image_path
+
 with gr.Blocks(theme="soft") as demo:
     # Block for "SNAPSCRIBE" centered at the top
     with gr.Row():
@@ -106,34 +85,38 @@ with gr.Blocks(theme="soft") as demo:
             gr.Markdown("<h2 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>Developed using RealVisXL_V5.0_Lightning model with ❤ by Aklavya</h2>")
     
     with gr.Row():
-        with gr.Column(scale=3):
-            prompt = gr.Textbox(
-                label="Input Prompt",
-                placeholder="Describe the image you want to create",
-                lines=2,
-            )
-            run_button = gr.Button("Generate Image")
-
-            # Example prompts box
-            example_prompts = gr.Textbox(
-                value="Dew-covered spider web in morning sunlight, with blurred greenery\n"
-                      "--------------------------------------------\n"
-                      "Glass of cold water with ice cubes and condensation on a wooden table\n"
-                      "--------------------------------------------\n"
-                      "Coffee cup with latte art, steam rising, and morning sunlight\n"
-                      "--------------------------------------------\n"
-                      "Autumn forest with golden leaves, sunlight through trees, and a breeze",
-                lines=5,
-                label="Sample Inputs",
-                interactive=False,
-            )
-        with gr.Column(scale=7):
-            result_image = gr.Image(
-                label="Generated Image",
-                type="filepath",
-                elem_id="output_image",
-                height=600  # Increased the height by 100%
-            )
+    with gr.Column(scale=3):
+        prompt = gr.Textbox(
+            label="Input Prompt",
+            placeholder="Describe the image you want to create",
+            lines=2,
+        )
+        run_button = gr.Button("Generate Image")
+
+        # Example prompts box with Copy to Clipboard button
+        example_prompts_text = (
+            "Dew-covered spider web in morning sunlight, with blurred greenery\n"
+            "--------------------------------------------\n"
+            "Glass of cold water with ice cubes and condensation on a wooden table\n"
+            "--------------------------------------------\n"
+            "Coffee cup with latte art, steam rising, and morning sunlight\n"
+            "--------------------------------------------\n"
+            "Autumn forest with golden leaves, sunlight through trees, and a breeze"
+        )
+
+        example_prompts = gr.Textbox(
+            value=example_prompts_text,
+            lines=5,
+            label="Sample Inputs",
+            interactive=False,
+        )
+    with gr.Column(scale=7):
+        result_image = gr.Image(
+            label="Generated Image",
+            type="filepath",
+            elem_id="output_image",
+            height=1024  # Increased the height by 100%
+        )
 
     run_button.click(
         fn=generate,
@@ -141,4 +124,25 @@ with gr.Blocks(theme="soft") as demo:
         outputs=[result_image],
     )
 
+    # # Footer with custom style and text
+    # gr.Markdown("""
+    # <style>
+    # .footer {
+    #     position: relative;
+    #     left: 0;
+    #     bottom: 0;
+    #     width: 100%;
+    #     background-color: white;
+    #     color: black;
+    #     text-align: center;
+    # }
+    # #output_image {
+    #     height: calc(100% + 512px); /* Dynamically increases height by 100% */
+    # }
+    # </style>
+    # <div class="footer">
+    # <p>Developed with ❤ by Aklavya(Bucky)</p>
+    # </div>
+    # """)
+
 demo.launch()