api

main.py

@@ -1,7 +1,11 @@
 import os
-from flask import Flask, jsonify, request
+import io
+import torch
+from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
-from transformers import pipeline
+from diffusers import AutoPipelineForImage2Image
+from diffusers.utils import load_image, make_image_grid
+from PIL import Image
 
 # Set the cache directory to a location with write permissions
 os.environ['TRANSFORMERS_CACHE'] = './transformers_cache'
@@ -10,20 +14,49 @@ os.environ['HF_HOME'] = './transformers_cache'
 app = Flask(__name__)
 CORS(app)
 
-# Load the sentiment-analysis pipeline from Hugging Face
-sentiment_analyzer = pipeline('sentiment-analysis')
+# Load the image-to-image pipeline from Hugging Face
+pipe = AutoPipelineForImage2Image.from_pretrained("RunDiffusion/Juggernaut-X-v10", torch_dtype=torch.float16).to("cuda")
+pipe.enable_xformers_memory_efficient_attention()
+pipe.enable_vae_tiling()  # Improve performance on large images
+pipe.enable_vae_slicing()  # Improve performance on large batches
 
 @app.route('/')
 def hello():
     return {"Goes Wrong": "Keeping it real"}
 
-@app.route('/analyze', methods=['GET'])
-def analyze():
-    sentence = request.args.get('sentence', default='hey there', type=str)
-    result = sentiment_analyzer(sentence)
-    print(result)
-    return jsonify(result)
+@app.route('/generate', methods=['POST'])
+def generate():
+    if 'image' not in request.files:
+        return jsonify({"error": "No image file provided"}), 400
 
+    image_file = request.files['image']
+    prompt = request.form.get('prompt', 'fleece hoodie, front zip, abstract pattern, GAP logo, high quality, photo')
+    negative_prompt = request.form.get('negative_prompt', 'low quality, bad quality, sketches, hanger')
+    guidance_scale = float(request.form.get('guidance_scale', 7))
+    num_images = int(request.form.get('num_images', 2))
+
+    sketch = Image.open(image_file)
+    
+    with torch.inference_mode():
+        images = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=sketch,
+            num_inference_steps=35,
+            guidance_scale=guidance_scale,
+            strength=0.5,
+            generator=torch.manual_seed(69),
+            num_images_per_prompt=num_images,
+        ).images
+
+    grid = make_image_grid(images, rows=1, cols=num_images)
+    
+    # Save the generated grid to a BytesIO object
+    img_byte_arr = io.BytesIO()
+    grid.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+
+    return send_file(img_byte_arr, mimetype='image/png')
 
 if __name__ == '__main__':
     app.run(debug=True)

requirements.txt

@@ -14,4 +14,8 @@ onnxruntime
 onnx
 pycocotools
 transformers
-torch
+torch
+diffusers
+PIL
+cv2
+numpy