| import gradio as gr |
| import torch |
| from PIL import Image |
| from run import StableRemix, run_remixing |
|
|
| pipe = StableRemix.from_pretrained( |
| "stabilityai/stable-diffusion-2-1-unclip", |
| torch_dtype=torch.float16, |
| variant="fp16" |
| ) |
| pipe = pipe.to('cuda') |
| pipe.enable_attention_slicing() |
|
|
| print('pipe loaded') |
|
|
|
|
| def remix(image1, image2, alpha): |
| |
|
|
| |
| images = run_remixing(pipe, image1, image2, [alpha]) |
| return images[0] |
| for idx, image in enumerate(images): |
| path = args.save_dir / f'remix_{idx}.png' |
| print('Saving remix to', path) |
| image.save(path) |
|
|
|
|
| demo = gr.Interface( |
| fn=remix, inputs=[gr.Image(image_mode='RGB', shape=[512, 512]), gr.Image(image_mode='RGB', shape=[512, 512]), gr.Slider(0.0, 1.0, 0.6)], outputs="image") |
| demo.launch() |
|
|
