Spaces:
Paused
Paused
| import spaces | |
| import gradio as gr | |
| import torch | |
| from diffusers import LCMScheduler, AutoPipelineForText2Image | |
| from diffusers import AutoPipelineForInpainting, LCMScheduler | |
| from diffusers import DiffusionPipeline, LCMScheduler | |
| from PIL import Image, ImageEnhance | |
| import io | |
| pipe = DiffusionPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-xl-base-1.0", | |
| variant="fp16", | |
| torch_dtype=torch.float32 | |
| ).to("cuda") | |
| # set scheduler | |
| pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) | |
| # Load and fuse lcm lora | |
| pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm") | |
| pipe.load_lora_weights("Pclanglais/wiki-model", weight_name="pytorch_lora_weights.safetensors", adapter_name="mickey") | |
| # Combine LoRAs | |
| pipe.set_adapters(["lcm", "mickey"], adapter_weights=[1.0, 1.0]) | |
| pipe.fuse_lora() | |
| def generate_image(prompt, num_inference_steps, guidance_scale): | |
| model_id = "stabilityai/stable-diffusion-xl-base-1.0" | |
| adapter_id = "latent-consistency/lcm-lora-sdxl" | |
| pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float32, variant="fp16") | |
| pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) | |
| pipe.to("cuda") | |
| # Load and fuse lcm lora | |
| pipe.load_lora_weights(adapter_id) | |
| pipe.fuse_lora() | |
| # Generate the image | |
| image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0] | |
| return image | |
| def inpaint_image(prompt, init_image, mask_image, num_inference_steps, guidance_scale): | |
| pipe = AutoPipelineForInpainting.from_pretrained( | |
| "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", | |
| torch_dtype=torch.float32, | |
| variant="fp16", | |
| ).to("cuda") | |
| pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) | |
| pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl") | |
| pipe.fuse_lora() | |
| if init_image is not None: | |
| init_image_path = init_image.name # Get the file path | |
| init_image = Image.open(init_image_path).resize((1024, 1024)) | |
| else: | |
| raise ValueError("Initial image not provided or invalid") | |
| if mask_image is not None: | |
| mask_image_path = mask_image.name # Get the file path | |
| mask_image = Image.open(mask_image_path).resize((1024, 1024)) | |
| else: | |
| raise ValueError("Mask image not provided or invalid") | |
| # Generate the inpainted image | |
| generator = torch.manual_seed(42) | |
| image = pipe( | |
| prompt=prompt, | |
| image=init_image, | |
| mask_image=mask_image, | |
| generator=generator, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| ).images[0] | |
| return image | |
| def generate_image_with_adapter(prompt, num_inference_steps, guidance_scale): | |
| generator = torch.manual_seed(0) | |
| # Generate the image | |
| image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator).images[0] | |
| return image | |
| with gr.Blocks(gr.themes.Soft()) as demo: | |
| with gr.Row(): | |
| image_output = gr.Image(label="Generated Image") | |
| with gr.Row(): | |
| with gr.Accordion(label="Wiki-Mouse Image Generation"): | |
| adapter_prompt_input = gr.Textbox(label="Prompt", placeholder="papercut, a cute fox") | |
| adapter_steps_input = gr.Slider(minimum=1, maximum=10, label="Inference Steps", value=4) | |
| adapter_guidance_input = gr.Slider(minimum=0, maximum=2, label="Guidance Scale", value=1) | |
| adapter_generate_button = gr.Button("Generate Image with Adapter") | |
| adapter_generate_button.click( | |
| generate_image_with_adapter, | |
| inputs=[adapter_prompt_input, adapter_steps_input, adapter_guidance_input], | |
| outputs=image_output | |
| ) | |
| demo.launch() | |