from prodiapy import Custom from prodiapy.util import load from PIL import Image from threading import Thread from utils import image_to_base64 import gradio as gr import gradio_user_history as gr_user_history pipe = Custom(os.getenv("PRODIA_API_KEY")) def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, batch_count, profile: gr.OAuthProfile | None): total_images = [] threads = [] def generate_one_image(): result = pipe.create( "/sd/generate", prompt=prompt, negative_prompt=negative_prompt, model=model, steps=steps, cfg_scale=cfg_scale, sampler=sampler, width=width, height=height, seed=seed ) job = pipe.wait_for(result) total_images.append(job['imageUrl']) for x in range(batch_count): t = Thread(target=generate_one_image) threads.append(t) t.start() for t in threads: t.join() for image in total_images: gr_user_history.save_image(label=prompt, image=Image.open(load(image)), profile=profile) return gr.update(value=total_images, preview=False) def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, batch_count): if input_image is None: return total_images = [] threads = [] def generate_one_image(): result = pipe.create( "/sd/transform", imageData=image_to_base64(input_image), denoising_strength=denoising, prompt=prompt, negative_prompt=negative_prompt, model=model, steps=steps, cfg_scale=cfg_scale, sampler=sampler, width=width, height=height, seed=seed ) job = pipe.wait_for(result) total_images.append(job['imageUrl']) for x in range(batch_count): t = Thread(target=generate_one_image) threads.append(t) t.start() for t in threads: t.join() return gr.update(value=total_images, preview=False) def upscale(image, scale, profile: gr.OAuthProfile | None): if image is None: return job = pipe.create( '/upscale', imageData=image_to_base64(image), resize=scale ) image = pipe.wait_for(job)['imageUrl'] gr_user_history.save_image(label=f'upscale by {scale}', image=Image.open(load(image)), profile=profile) return image