Create app.py

app.py

@@ -0,0 +1,119 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionXLPipeline
+# from PIL import Image
+import concurrent.futures
+# import time
+
+# Load the models with appropriate pipeline
+model_cache = {}
+
+
+def load_model(model_name):
+    # Check if the model is already cached to avoid reloading every time
+    if model_name in model_cache:
+        return model_cache[model_name]
+
+    print(f"Loading model: {model_name}")
+    try:
+        if model_name == "SG161222/RealVisXL_V5.0_Lightning":
+            model = StableDiffusionXLPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
+        else:
+            model = StableDiffusionXLPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
+
+        model.to("cuda")
+        model_cache[model_name] = model  # Cache the model for future use
+        print("Model loaded successfully.")
+        return model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
+
+
+# Function to generate the image with a timeout
+def generate_image_with_timeout(prompt, model_name):
+    # Set a timeout for the image generation (120 seconds)
+    timeout = 180  # 120 seconds
+
+    try:
+        # Use ThreadPoolExecutor to handle the timeout
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future = executor.submit(generate_image, prompt, model_name)
+            return future.result(timeout=timeout)  # Will raise TimeoutError if the process exceeds timeout
+
+    except concurrent.futures.TimeoutError:
+        return "Error: The image generation timed out. Please try again."
+
+
+# Function to generate the image
+def generate_image(prompt, model_name):
+    # Load the appropriate model
+    model = load_model(model_name)
+
+    if model is None:
+        return "Error loading the model."
+
+    try:
+        # Generate the image from the prompt
+        with torch.no_grad():
+            image = model(prompt).images[0]
+        return image
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return "Error generating the image."
+
+
+# Define the Gradio interface using gr.Blocks
+def create_gradio_interface():
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        # Add a heading at the top of the screen
+        gr.Markdown("""
+            <h1 style="
+                text-align: center;
+                color: white;
+                font-weight: bold;
+                text-transform: uppercase;
+                text-decoration: underline;
+                margin-top: 30px;
+                font-family: 'Arial', sans-serif;
+                background: linear-gradient(45deg, #ff6b6b, #f06595);
+                padding: 10px 20px;
+                border-radius: 15px;
+                box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.3);
+            ">
+                SNAPSCRIBE
+            </h1>
+        """)
+
+        # Create a Row for input-output arrangement
+        with gr.Row():
+            # Create a column for the left input section
+            with gr.Column(scale=0.3, min_width=300):
+                prompt_input = gr.Textbox(label="Enter your prompt here", placeholder="e.g., A futuristic city skyline")
+                model_input = gr.Dropdown(
+                    choices=["SG161222/RealVisXL_V5.0_Lightning", "SG161222/RealVisXL_V4.0_Lightning"],
+                    label="Choose model",
+                    value="SG161222/RealVisXL_V5.0_Lightning"
+                )
+                submit_button = gr.Button("Generate Image")
+
+            # Create a column for the right output image section with reduced height (20% smaller)
+            with gr.Column(scale=0.7, min_width=600):
+                output_image = gr.Image(label="Generated Image", height=640)  # Reduced height by 20%
+
+        # Bind the function to the Gradio interface
+        submit_button.click(fn=generate_image_with_timeout, inputs=[prompt_input, model_input], outputs=output_image)
+
+        # Add footer with a styled span for text in the footer
+        gr.Markdown("""
+            <div style="position: relative; left: 0; bottom: 0; width: 100%; background-color: #0B0F19; color: white; text-align: center; padding: 10px 0;">
+                <p>Developed with ❤ by Aklavya (Bucky)</p>
+            </div>
+        """)
+
+    # Launch the interface
+    demo.launch(share=True)
+
+
+# Launch the Gradio interface
+create_gradio_interface()