Update app.py

app.py

@@ -3,18 +3,19 @@ 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": (
-            "hyper-realistic image of {prompt}. lifelike, authentic, natural colors, true-to-life details, realistic lighting, immersive, highly detailed",
-            "unrealistic, high resolution, artificial, over-saturated, distorted, fake"
-
+            "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+            "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
         ),
         "Style Zero": ("{prompt}", ""),
     }
@@ -23,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(
@@ -34,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,
@@ -58,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():
@@ -85,6 +105,7 @@ with gr.Blocks(theme="soft") as demo:
             gr.Markdown("<h1 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>SNAPSCRIBE</h1>")
             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",
@@ -92,25 +113,20 @@ with gr.Blocks(theme="soft") as demo:
                 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 box
             example_prompts = gr.Textbox(
-                value=example_prompts_text,
+                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",
@@ -118,10 +134,6 @@ with gr.Blocks(theme="soft") as demo:
                 elem_id="output_image",
                 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,
@@ -129,25 +141,4 @@ with gr.Blocks(theme="soft") as demo:
         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()
+demo.launch()