Spaces:
Runtime error
Runtime error
File size: 2,157 Bytes
781570f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image, make_image_grid
from PIL import Image
import torch
import io
# Set environment variable to avoid fragmentation
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
# Clear any unused GPU memory
torch.cuda.empty_cache()
app = Flask(__name__)
CORS(app)
# 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('/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)
|