Spaces:
Runtime error
Runtime error
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 | |
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() | |