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)