🎉 | Project added to Git

main.py

@@ -0,0 +1,58 @@
+from flask import Flask, request, jsonify
+from diffusers import DDPMScheduler, DiffusionPipeline
+import torch
+import base64
+from io import BytesIO
+from PIL import Image
+
+app = Flask(__name__)
+
+scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", scheduler=scheduler)
+
+print("Loading UNet weights...")
+ckpt_path = './models/model.bin'
+try:
+    unet_state_dict = torch.load(ckpt_path)
+    pipe.unet.load_state_dict(unet_state_dict)
+    if torch.cuda.is_available():
+        pipe = pipe.to("cuda")
+        pipe.to(torch.float16)
+        print("UNet weights loaded successfully and model moved to CUDA.")
+    else:
+        print("CUDA is not available. Running on CPU (will be slow).")
+        pipe = pipe.to("cpu")
+        pipe.to(torch.float32)
+except FileNotFoundError:
+    print(f"Error: Model not found at {ckpt_path}. Please ensure the file exists.")
+    exit()
+except Exception as e:
+    print(f"An error occurred while loading UNet weights: {e}")
+    exit()
+
+@app.route('/api/generate', methods=['POST'])
+def generate_image():
+    prompt = request.json.get('prompt')
+    if not prompt:
+        return jsonify({"error": "Prompt is required in the request body."}), 400
+
+    print(f"Generating image for prompt: '{prompt}'...")
+    try:
+        image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[999]).images[0]
+        image = image.resize((128, 128), Image.LANCZOS)
+
+        # Convert image to base64
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG") # You can choose JPEG for smaller size, but PNG is lossless
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        return jsonify({"image": img_str})
+
+    except Exception as e:
+        print(f"Error during image generation: {e}")
+        return jsonify({"error": "An error occurred during image generation."}), 500
+
+if __name__ == '__main__':
+    # Make sure to run this with 'python app.py'
+    # It will be accessible at http://127.0.0.1:5000/
+    app.run(host='127.0.0.1', port=5000, debug=True) # debug=True is good for development