Update app.py

app.py

@@ -1,35 +1,44 @@
 import gradio as gr
 
-
+# Define the available models
 models = {
     "Flux Lora": "models/prashanth970/flux-lora-uncensored",
     "TrioHMH Flux": "models/DiegoJR1973/NSFW-TrioHMH-Flux",
     "Master": "models/pimpilikipilapi1/NSFW_master"
 }
 
+# Function to generate an image from text using the selected model
 def generate_image(text, model_name):
-    model_path = models[model_name]
+    model_path = models[model_name]  # Get the path of the selected model
     print(f"Fetching model from: {model_path}")
     
     try:
-        model = gr.load(model_path)
-        result_image = model(text)
-        return result_image
+        # Dynamically load the model based on the selected model path
+        model = gr.load(model_path)  # Ensure this is the correct method to load your model
+        result_image = model(text)  # Generate the image from the input text
+        
+        # Ensure the result is in a proper image format
+        if isinstance(result_image, str):  # if model returns a path to the image
+            return gr.Image(result_image)
+        elif isinstance(result_image, bytes):  # if model returns raw image bytes
+            return gr.Image(value=result_image)
+        else:
+            return result_image
     except Exception as e:
         print(f"Error loading model: {e}")
         return None
 
-
+# Gradio Interface setup
 interface = gr.Interface(
     fn=generate_image,
     inputs=[
-        gr.Textbox(label="Type here your imagination:", placeholder="Type your description here..."),
-        gr.Dropdown(label="Select Model", choices=list(models.keys()), value="Flux Lora")
+        gr.Textbox(label="Type here your imagination:", placeholder="Type your description here..."),  # Textbox for input
+        gr.Dropdown(label="Select Model", choices=list(models.keys()), value="Flux Lora")  # Dropdown for selecting the model
     ],
-    outputs=gr.Image(label="Generated Image"),
-    theme="NoCrypt/miku",
+    outputs=gr.Image(label="Generated Image"),  # Image output
+    theme="NoCrypt/miku",  # Set theme for the interface
     description="Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.",
 )
 
-
+# Launch the Gradio interface
 interface.launch()