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  # Placeholder; adjust based on InstantMesh docs
+from diffusers import StableDiffusionPipeline
 from PIL import Image
 import logging
 
@@ -12,17 +12,13 @@ logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
 
-# Load the model once at startup (on CPU) with token from environment
+# Load the model once at startup (on CPU)
 try:
-    logger.info("Loading TencentARC InstantMesh pipeline...")
-    token = os.getenv("HF_TOKEN")
-    if not token:
-        raise ValueError("HF_TOKEN environment variable not set")
+    logger.info("Loading runwayml/stable-diffusion-v1-5 pipeline...")
     pipe = StableDiffusionPipeline.from_pretrained(
-        "TencentARC/InstantMesh",
+        "runwayml/stable-diffusion-v1-5",
         torch_dtype=torch.float32,
         cache_dir="/tmp/hf_home",
-        token=token,  # Reintroduce token authentication
     )
     pipe.to("cpu")
     logger.info("=== Application Startup at CPU mode =====")
@@ -37,7 +33,7 @@ def pil_to_base64(image):
 
 @app.route("/")
 def home():
-    return "TencentARC InstantMesh CPU API is running!"
+    return "Stable Diffusion CPU API is running!"
 
 @app.route("/generate", methods=["POST"])
 def generate():
@@ -46,30 +42,20 @@ def generate():
 
     try:
         data = request.get_json()
-        image_data = data.get("image")
+        prompt = data.get("prompt")  # Use text prompt instead of image
 
-        if not image_data:
-            return jsonify({"error": "No image provided"}), 400
+        if not prompt:
+            return jsonify({"error": "No prompt provided"}), 400
 
-        if image_data.startswith("data:image"):
-            image_data = image_data.split(",")[1]
+        logger.info("Generating image with pipeline...")
+        result = pipe(prompt)  # Generate image from text
+        image = result.images[0]  # Get the first generated image
 
-        image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
-
-        logger.info("Processing image with pipeline...")
-        result = pipe(image)  # Adjust based on InstantMesh documentation
-        output_mesh = result.mesh  # Hypothetical; check InstantMesh output format
-
-        output_path = "/tmp/output.glb"
-        output_mesh.save(output_path)
-        with open(output_path, "rb") as f:
-            mesh_data = base64.b64encode(f.read()).decode("utf-8")
-
-        logger.info("Mesh processed successfully")
-        return jsonify({"mesh": f"data:model/gltf-binary;base64,{mesh_data}"})
+        logger.info("Image generated successfully")
+        return jsonify({"image": f"data:image/png;base64,{pil_to_base64(image)}"})
 
     except Exception as e:
-        logger.error(f"Error generating mesh: {e}", exc_info=True)
+        logger.error(f"Error generating image: {e}", exc_info=True)
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":