Update app.py

app.py

@@ -1,8 +1,6 @@
 import gradio as gr
 import time
 import logging
-import os
-from PIL import Image
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -11,18 +9,17 @@ logger = logging.getLogger(__name__)
 def load_model(model_path):
     logger.info(f"Attempting to load model from {model_path}...")
     try:
-        if os.path.exists(model_path):
-            demo = gr.load(model_path)
-            logger.info(f"Successfully loaded demo from {model_path}.")
-            return demo
-        else:
-            raise FileNotFoundError(f"Model file does not exist at {model_path}")
+        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
 
-def generate_image(prompt, model_path, max_retries=5):
-    demo = load_model(model_path)
+# 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."
 
@@ -33,6 +30,7 @@ def generate_image(prompt, model_path, max_retries=5):
             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)}")
@@ -41,50 +39,20 @@ def generate_image(prompt, model_path, max_retries=5):
             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"
 
-def load_image_from_file(file_path):
-    try:
-        if os.path.exists(file_path):
-            logger.info(f"Loading image from {file_path}")
-            return Image.open(file_path), "Image loaded successfully"
-        else:
-            return None, "Error: File does not exist"
-    except Exception as e:
-        logger.error(f"Failed to load image from {file_path}: {e}")
-        return None, f"Error loading image: {e}"
-
-# Define the Gradio interface
-with gr.Blocks() as demo:
-    with gr.Row():
-        # Left side: prompt-based image generation
-        with gr.Column():
-            prompt = gr.Textbox(label="Enter your prompt", lines=2)
-            model_path = gr.FileExplorer(label="Select Model Directory")
-            generate_btn = gr.Button("Generate Image")
-            generated_image = gr.Image(label="Generated Image")
-            output_info = gr.Textbox(label="Output Information", lines=5)
-
-        # Right side: file loading
-        with gr.Column():
-            image_path = gr.FileExplorer(label="Select Image File")
-            load_image_btn = gr.Button("Load Image")
-            loaded_image = gr.Image(label="Loaded Image")
-
-    # Logic for generating images
-    generate_btn.click(
-        fn=generate_image, 
-        inputs=[prompt, model_path], 
-        outputs=[generated_image, output_info]
-    )
-    
-    # Logic for loading images from the file system
-    load_image_btn.click(
-        fn=load_image_from_file, 
-        inputs=[image_path], 
-        outputs=[loaded_image, output_info]
-    )
+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...")
-demo.launch()
+iface.launch()