Update app.py

app.py

@@ -1,44 +1,42 @@
 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]  # Get the path of the selected model
+    model_path = models[model_name]
     print(f"Fetching model from: {model_path}")
     
     try:
-        # 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
+        model = gr.load(model_path)
+        result_image = model(text)
+        if isinstance(result_image, str):
+            return gr.Image(value=result_image)
+        elif isinstance(result_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..."),  # Textbox for input
-        gr.Dropdown(label="Select Model", choices=list(models.keys()), value="Flux Lora")  # Dropdown for selecting the model
+        gr.Textbox(label="Type here your imagination:", placeholder="Type your description here..."),
+        gr.Dropdown(label="Select Model", choices=list(models.keys()), value="Flux Lora")
     ],
-    outputs=gr.Image(label="Generated Image"),  # Image output
-    theme="NoCrypt/miku",  # Set theme for the interface
+    outputs=gr.Image(label="Generated Image"),
+    theme="NoCrypt/miku",
     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()
+
+interface.launch()