from flask import Flask, request, jsonify
from PIL import Image
import torch
from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
from diffusers.utils import load_image

app = Flask(__name__)

controlnet = ControlNetModel.from_pretrained("DionTimmer/controlnet_qrcode-control_v1p_sd15",
                                             torch_dtype=torch.float16, local_files_only=False)

pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    safety_checker=None,
    torch_dtype=torch.float16
)

pipe.enable_xformers_memory_efficient_attention()
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

def resize_for_condition_image(input_image: Image, resolution: int):
    input_image = input_image.convert("RGB")
    W, H = input_image.size
    k = float(resolution) / min(H, W)
    H *= k
    W *= k
    H = int(round(H / 64.0)) * 64
    W = int(round(W / 64.0)) * 64
    img = input_image.resize((W, H), resample=Image.LANCZOS)
    return img

@app.route('/generate_image', methods=['POST'])
def generate_image():
    # Get input parameters from the request
    prompt = request.json.get('prompt')
    negative_prompt = request.json.get('negative_prompt')
    image_url = request.json.get('image_url')
    control_image_url = request.json.get('control_image_url')

    # Load the images from URLs
    source_image = load_image(image_url)
    init_image = load_image(control_image_url)
    
    # Resize images for conditioning
    condition_image = resize_for_condition_image(source_image, 768)
    init_image = resize_for_condition_image(init_image, 768)
    
    # Generate the image using the pipeline
    generator = torch.manual_seed(123121231)
    image = pipe(prompt=prompt,
                 negative_prompt=negative_prompt, 
                 image=init_image,
                 control_image=condition_image,
                 width=768,
                 height=768,
                 guidance_scale=20,
                 controlnet_conditioning_scale=1.5,
                 generator=generator,
                 strength=0.9, 
                 num_inference_steps=150)
    
    # Return the generated image
    return jsonify({'image': image.images[0]})

if __name__ == '__main__':
    app.run()