Update app.py

app.py

@@ -3,7 +3,7 @@ import base64
 import torch
 import os
 from flask import Flask, request, jsonify
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline  # Adjust if needed based on ldm
 from PIL import Image
 import logging
 
@@ -14,9 +14,9 @@ app = Flask(__name__)
 
 # Load the model once at startup (on CPU)
 try:
-    logger.info("Loading runwayml/stable-diffusion-v1-5 pipeline...")
+    logger.info("Loading CompVis/ldm-stable-diffusion-v1 pipeline...")
     pipe = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
+        "CompVis/ldm-stable-diffusion-v1",
         torch_dtype=torch.float32,
         cache_dir="/tmp/hf_home",
     )
@@ -33,7 +33,7 @@ def pil_to_base64(image):
 
 @app.route("/")
 def home():
-    return "Stable Diffusion CPU API is running!"
+    return "CompVis ldm-stable-diffusion-v1 CPU API is running!"
 
 @app.route("/generate", methods=["POST"])
 def generate():
@@ -42,7 +42,7 @@ def generate():
 
     try:
         data = request.get_json()
-        prompt = data.get("prompt")  # Use text prompt instead of image
+        prompt = data.get("prompt")  # Use text prompt for 2D image generation
 
         if not prompt:
             return jsonify({"error": "No prompt provided"}), 400