Create app.py

app.py

@@ -0,0 +1,57 @@
+import gradio as gr
+import time
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+def load_model(model_path):
+    logger.info(f"Attempting to load model from {model_path}...")
+    try:
+        demo = gr.load(model_path)
+        logger.info(f"Successfully loaded demo from {model_path}.")
+        return demo
+    except Exception as e:
+        logger.error(f"Failed to load demo from {model_path}: {e}")
+        return None
+
+# Load the demo
+demo = load_model("models/votepurchase/NSFW-gen-v2")
+
+def generate_image(prompt, max_retries=5):
+    # if demo is None:
+    #     return None, "Error: Demo failed to load."
+
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Attempt {attempt + 1}: Starting image generation with prompt: {prompt}")
+            start_time = time.time()
+            result = demo(prompt)
+            generation_time = time.time() - start_time
+            logger.info(f"Generation completed in {generation_time:.2f} seconds")
+            return result, f"Image generated successfully\nGeneration time: {generation_time:.2f} seconds"
+        except Exception as e:
+            logger.error(f"Error on attempt {attempt + 1}: {str(e)}")
+            if attempt < max_retries - 1:
+                logger.info(f"Retrying... ({attempt + 2}/{max_retries})")
+            else:
+                logger.error("Max retries reached. Giving up.")
+                return None, f"Failed to generate image after {max_retries} attempts. Last error: {str(e)}"
+    
+    # This line should never be reached, but just in case:
+    return None, "Unexpected error occurred"
+
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Enter your prompt"),
+    outputs=[
+        gr.Image(label="Generated Image"),
+        gr.Textbox(label="Output Information", lines=5)
+    ],
+    title="Text-to-Image Generation",
+    description="Enter a prompt to generate an image. The system will automatically retry up to 5 times if an error occurs."
+)
+
+logger.info("Launching the Gradio interface...")
+iface.launch()