Update app.py

app.py

@@ -7,6 +7,43 @@ logging.basicConfig(
 )
 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/John6666/lustify-sdxl-nsfwsfw-v2-sdxl")
+
+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")
+            logger.info(result)
+            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"
+
 # Set the root directory as a reusable variable
 ROOT_DIR = "./tmp"
 
@@ -79,6 +116,5 @@ with gr.Blocks() as demo:
     )
 
 logger.info("Launching Gradio app")
-
-if __name__ == "__main__":
-    demo.launch(debug=True)
+demo.launch()
+logger.info("Gradio app launched")