Update app.py

app.py

@@ -5,7 +5,18 @@ import torch
 import numpy as np
 from PIL import Image, ImageEnhance
 import spaces
-
+from pymongo import MongoClient
+from pymongo.errors import ConnectionError
+
+# MongoDB connection
+mongo_client = None
+try:
+    mongo_client = MongoClient("mongodb+srv://skandanv:Skandanv00031@cluster1.aquia.mongodb.net/?retryWrites=true&w=majority&appName=cluster1")
+    db = mongo_client['minecraft_skin_generator']  # Replace with your database name
+    collection = db['generated_skins']  # Collection to store generated skins
+    connection_message = "Connected to MineSkin Server"
+except ConnectionError:
+    connection_message = "Failed to connect to MineSkin Server"
 
 if torch.cuda.is_available():
     device = "cuda"
@@ -14,19 +25,14 @@ else:
     device = "cpu"
     print("Using CPU")
 
-
 MAX_SEED = np.iinfo(np.int32).max
 
-
 subprocess.run(["git", "clone", "https://github.com/Nick088Official/Stable_Diffusion_Finetuned_Minecraft_Skin_Generator.git"])
-
 os.chdir("Stable_Diffusion_Finetuned_Minecraft_Skin_Generator")
 
-
 @spaces.GPU(duration=75)
 def run_inference(prompt, stable_diffusion_model, num_inference_steps, guidance_scale, model_precision_type, seed, filename, model_3d, verbose):
-
-    # inference
+    # Inference
     if stable_diffusion_model == '2':
         sd_model = "minecraft-skins"
     elif stable_diffusion_model == 'xl':
@@ -36,12 +42,34 @@ def run_inference(prompt, stable_diffusion_model, num_inference_steps, guidance_
     
     os.system(inference_command)
 
-    # view it also as 3d model or not output
-    if model_3d:
-        return os.path.join(f"output_minecraft_skins/{filename}"), os.path.join(f"output_minecraft_skins/{filename}_3d_model.glb")
-    else:
-        return os.path.join(f"output_minecraft_skins/{filename}"), None
-
+    # File paths for generated assets
+    image_path = os.path.join(f"output_minecraft_skins/{filename}")
+    model_path = os.path.join(f"output_minecraft_skins/{filename}_3d_model.glb") if model_3d else None
+
+    # Prepare data for MongoDB
+    skin_data = {
+        'prompt': prompt,
+        'filename': filename,
+        'image_path': image_path,
+        'model_path': model_path,
+        'num_inference_steps': num_inference_steps,
+        'guidance_scale': guidance_scale,
+        'model_precision_type': model_precision_type,
+        'seed': seed,
+        'model_3d': model_3d,
+        'verbose': verbose
+    }
+
+    # Insert generated skin data into MongoDB and show alert if successful
+    try:
+        collection.insert_one(skin_data)
+        success_message = "The Skin has been pushed to MineSkin Server"
+        alert_type = "success"  # Gradio Alert type for success
+    except Exception as e:
+        success_message = f"Failed to push skin to database: {e}"
+        alert_type = "error"
+
+    return image_path, model_path, success_message, alert_type
 
 # Define Gradio UI components
 prompt = gr.Textbox(label="Your Prompt", info="What the Minecraft Skin should look like")
@@ -50,12 +78,15 @@ num_inference_steps = gr.Slider(label="Number of Inference Steps", info="The num
 guidance_scale = gr.Slider(label="Guidance Scale", info="Controls how much the image generation process follows the text prompt. Higher values make the image stick more closely to the input text.", minimum=0.0, maximum=10.0, value=7.5, step=0.1)
 model_precision_type = gr.Dropdown(["fp16", "fp32"], value="fp16", label="Model Precision Type", info="The precision type to load the model, like fp16 which is faster, or fp32 which is more precise but more resource consuming")
 seed = gr.Slider(value=42, minimum=0, maximum=MAX_SEED, step=1, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one")
-filename = gr.Textbox(label="Output Image Name", info="The name of the file of the output image skin, keep the.png", value="output-skin.png")
+filename = gr.Textbox(label="Output Image Name", info="The name of the file of the output image skin, keep the .png", value="output-skin.png")
 model_3d = gr.Checkbox(label="See as 3D Model too", info="View the generated skin as a 3D Model too", value=True)
 verbose = gr.Checkbox(label="Verbose Output", info="Produce more detailed output while running", value=False)
 
-
 # Create the Gradio interface
+output_image = gr.Image(label="Generated Minecraft Skin Image Asset", elem_classes="pixelated checkered")
+output_model = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model View of the Skin")
+output_message = gr.Alert()
+
 gr.Interface(
     fn=run_inference,
     inputs=[
@@ -70,10 +101,16 @@ gr.Interface(
         verbose
     ],
     outputs=[
-        gr.Image(label="Generated Minecraft Skin Image Asset", elem_classes="pixelated checkered"),
-        gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model View of the Skin")
+        output_image,
+        output_model,
+        output_message
     ],
     title="Minecraft Skin Generator",
-    description="Make AI generated Minecraft Skins by a Finetuned Stable Diffusion Version!<br>Github Repository & Model used: https://github.com/Nick088Official/Stable_Diffusion_Finetuned_Minecraft_Skin_Generator<br>Credits: [Monadical-SAS](https://github.com/Monadical-SAS/minecraft_skin_generator) (Creators of the model), [Nick088](https://linktr.ee/Nick088) (Improving usage of the model), daroche (helping me fix the 3d model texture isue), [Brottweiler](https://gist.github.com/Brottweiler/483d0856c6692ef70cf90bf1a85ce364)(script to fix the 3d model texture), [not-holar](https://huggingface.co/not-holar) (made the rendering of the image asset in the web ui look pixelated like minecraft and have a checkered background),[meew](https://huggingface.co/spaces/meeww/Minecraft_Skin_Generator/blob/main/models/player_model.glb) (Minecraft Player 3d model) <br> [![Discord](https://img.shields.io/discord/1198701940511617164?color=%23738ADB&label=Discord&style=for-the-badge)](https://discord.gg/AQsmBmgEPy)",
-    css=".pixelated {image-rendering: pixelated} .checkered img {background-image: url(\'data:image/svg+xml,<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"2\" height=\"2\" fill-opacity=\".15\"><rect x=\"1\" width=\"1\" height=\"1\"/><rect y=\"1\" width=\"1\" height=\"1\"/></svg>\');background-size: 16px;}"
-).launch(show_api=True, share=True)
+    description="Make AI generated Minecraft Skins by a Finetuned Stable Diffusion Version!<br>Github Repository & Model used: https://github.com/Nick088Official/Stable_Diffusion_Finetuned_Minecraft_Skin_Generator<br>Credits: [Monadical-SAS](https://github.com/Monadical-SAS/minecraft_skin_generator) (Creators of the model), [Nick088](https://linktr.ee/Nick088) (Improving usage of the model), daroche (helping me fix the 3d model texture issue), [Brottweiler](https://gist.github.com/Brottweiler/483d0856c6692ef70cf90bf1a85ce364)(script to fix the 3d model texture), [not-holar](https://huggingface.co/not-holar) (made the rendering of the image asset in the web ui look pixelated like minecraft and have a checkered background),[meew](https://huggingface.co/spaces/meeww/Minecraft_Skin_Generator/blob/main/models/player_model.glb) (Minecraft Player 3d model) <br> [![Discord](https://img.shields.io/discord/1198701940511617164?color=%23738ADB&label=Discord&style=for-the-badge)](https://discord.gg/AQsmBmgEPy)",
+    css=".pixelated {image-rendering: pixelated} .checkered img {background-image: url('data:image/svg+xml,<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"2\" height=\"2\" fill-opacity=\".15\"><rect x=\"1\" width=\"1\" height=\"1\"/><rect y=\"1\" width=\"1\" height=\"1\"/></svg>');background-size: 16px;}"
+).launch(show_api=True, share=True)
+
+# Show connection message as alert when the app starts
+if mongo_client:
+    gr.Interface().launch()
+    output_message.update(value=connection_message, type="info")