UPDATE refined

app.py

@@ -1,29 +1,70 @@
+import logging
+from typing import Any
 import gradio as gr
 from somnium import Somnium
 
-# define function to generate image
-def generate_image(prompt, style_id):
-    if prompt == "" or style_id == "":
-        raise gr.Error("Empty values")
+logging.basicConfig(
+    level=logging.INFO, 
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+try:
+    STYLES_DICT = Somnium.Styles()
+    STYLE_CHOICES = list(STYLES_DICT.keys())
+except Exception as e:
+    logging.error(f"Failed to load Somnium styles: {e}")
+    STYLES_DICT = {}
+    STYLE_CHOICES = []
+
+def generate_image(prompt: str, style_id: str) -> Any:
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a valid prompt.")
+    if not style_id or style_id not in STYLES_DICT:
+        raise gr.Error("Please select a valid style.")
+
     try:
-        styles = Somnium.Styles()
-        image = Somnium.Generate(prompt, styles[style_id])
-        image + "69" # Secret line
+        logging.info(f"Generating image with style: '{style_id}'")
+        image = Somnium.Generate(prompt, STYLES_DICT[style_id])
         return image
     except Exception as e:
-        raise gr.Error("Process failed or contains NSFW")
-
-# create interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(label="Enter Prompt:", max_lines=10),
-        gr.Dropdown(list((Somnium.Styles()).keys()), label="Select Style:")
-    ],
-    outputs=gr.Image(show_download_button=False, show_share_button=False),
-    allow_duplication=True,
-    title="Somnium Image Generator"
-)
+        logging.error(f"Generation failed: {e}", exc_info=True)
+        raise gr.Error("Generation failed. Please try a different prompt or check for NSFW content.")
+
+
+with gr.Blocks(title="Somnium Image Generator") as demo:
+    
+    gr.Markdown(
+        """
+        # 🌌 Somnium Image Generator
+        Enter a prompt and select a style to generate your image.
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(
+                label="Enter Prompt:", 
+                placeholder="A futuristic city at sunset...",
+                lines=5
+            )
+            style_dropdown = gr.Dropdown(
+                choices=STYLE_CHOICES, 
+                label="Select Style:",
+                value=STYLE_CHOICES[0] if STYLE_CHOICES else None
+            )
+            generate_btn = gr.Button("Generate Image", variant="primary")
+            
+        with gr.Column(scale=1):
+            image_output = gr.Image(
+                label="Output", 
+                buttons=["download"] 
+            )
+
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt_input, style_dropdown],
+        outputs=image_output
+    )
 
-# run the interface
-iface.launch()
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Soft())