Update app.py

app.py

@@ -3,12 +3,14 @@ import random
 import uuid
 import gradio as gr
 import numpy as np
-from PIL import Image
-import spaces
 import torch
+import concurrent.futures
+from PIL import Image
 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": (
@@ -22,6 +24,8 @@ 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(
@@ -33,19 +37,25 @@ 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
 
-@spaces.GPU(duration=60, enable_queue=True)
+
+# Main generation function with timeout
 def generate(
         prompt: str,
         seed: int = 1,
@@ -57,26 +67,37 @@ def generate(
 ):
     global 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
-    image_path = save_image(images[0])  # Saving the first generated image
-    return image_path
-
+    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
 with gr.Blocks(theme="soft") as demo:
     # Block for "SNAPSCRIBE" centered at the top
     with gr.Row():
@@ -94,7 +115,6 @@ with gr.Blocks(theme="soft") as demo:
             run_button = gr.Button("Generate Image")
 
             # Example prompts box
-            # gr.Markdown("<h3 style='text-align:center;'>Sample Inputs</h3>")
             example_prompts = gr.Textbox(
                 value="Dew-covered spider web in morning sunlight, with blurred greenery\n"
                       "--------------------------------------------\n"
@@ -115,34 +135,10 @@ with gr.Blocks(theme="soft") as demo:
                 height=600  # Increased the height by 100%
             )
 
-        # with gr.Column(scale=12):
-        #     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</h1>")
-
     run_button.click(
         fn=generate,
         inputs=[prompt],
         outputs=[result_image],
     )
 
-    # # Footer added to center-align the text
-    # gr.HTML("""
-    # <style>
-    #     .footer {
-    #         position: relative;
-    #         left: 0;
-    #         bottom: 0;
-    #         width: 100%;
-    #         background-color: white;
-    #         color: black;
-    #         text-align: center;
-    #         padding: 10px;
-    #         margin-top: 20px;
-    #         font-weight: bold;
-    #     }
-    # </style>
-    # <div class="footer">
-    #     <p>Developed using RealVisXL_V5.0_Lightning model with ❤ by Aklavya (Bucky)</p>
-    # </div>
-    # """)
-
 demo.launch()