switch to LCM demo

README.md

@@ -1,8 +1,8 @@
 ---
-title: SSD-1B + SDXL 1.0 & Refiner
+title: LCM-LoRa-SDXL Demo + Papercut
 emoji: 🏆
 colorFrom: indigo
-colorTo: pink
+colorTo: purple
 sdk: gradio
 sdk_version: 4.1.2
 app_file: app.py

app.py

@@ -1,68 +1,152 @@
 import spaces
-from diffusers import StableDiffusionXLPipeline
-from diffusers import DiffusionPipeline
-from pydantic import BaseModel
-from PIL import Image
 import gradio as gr
 import torch
-import uuid
+from diffusers import LCMScheduler, AutoPipelineForText2Image
+from diffusers import AutoPipelineForInpainting, LCMScheduler
+from diffusers import DiffusionPipeline, LCMScheduler
+from PIL import Image, ImageEnhance
 import io
-import os
 
+@spaces.GPU
+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"
 
-# Load your model
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "segmind/SSD-1B", 
-    torch_dtype=torch.float16, 
-    use_safetensors=True, 
-    variant="fp16"
-)
+    pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.to("cuda")
 
-pipe.to("cuda:0")
+    # 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.float16,
+        variant="fp16",
+    ).to("cuda")
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
+    pipe.fuse_lora()
 
-@spaces.GPU 
-def generate_and_save_image(prompt, negative_prompt=''):
-    # Generate image using the provided prompts
-    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
+    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")
 
-    # Generate a unique UUID for the filename
-    unique_id = str(uuid.uuid4())
-    image_path = f"generated_images/{unique_id}.jpeg"
+    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")
 
-    # Save generated image locally
-    os.makedirs('generated_images', exist_ok=True)
-    image.save(image_path, format='JPEG')
+    # 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 the path of the saved image to display in Gradio interface
-    return image_path
+    return image
 
+def generate_image_with_adapter(prompt, num_inference_steps, guidance_scale):
+    pipe = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    variant="fp16",
+    torch_dtype=torch.float16
+    ).to("cuda")
 
-# Start of the Gradio Blocks interface
-with gr.Blocks() as demo:
-    with gr.Column():
-        gr.Markdown("# Image Generation with SSD-1B")
-        gr.Markdown("Enter a prompt and (optionally) a negative prompt to generate an image.")
-        
-        # Input fields for positive and negative prompts
-        with gr.Row():
-            prompt1 = gr.Textbox(label="Enter prompt")
-            negative_prompt = gr.Textbox(label="Enter negative prompt (optional)")
-        
-        # Button for generating the image
-        generate_button1 = gr.Button("Generate Image")
+    # set scheduler
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 
-        # Output image display, set to a larger default size
-        output_image1 = gr.Image(label="Generated Image")
+    # Load and fuse lcm lora
+    pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm")
+    pipe.load_lora_weights("TheLastBen/Papercut_SDXL", weight_name="papercut.safetensors", adapter_name="papercut")
 
-        # Click event for the generate button
-        generate_button1.click(
-            generate_and_save_image, 
-            inputs=[prompt1, negative_prompt], 
-            outputs=output_image1
-        )
+    # Combine LoRAs
+    pipe.set_adapters(["lcm", "papercut"], adapter_weights=[1.0, 0.8])
+    pipe.fuse_lora()
+    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]
+    pipe.unfuse_lora()
+    return image
 
 
-# Launch the combined Gradio app
+def modify_image(image, brightness, contrast):
+    # Function to modify brightness and contrast
+    image = Image.open(io.BytesIO(image))
+    enhancer = ImageEnhance.Brightness(image)
+    image = enhancer.enhance(brightness)
+    enhancer = ImageEnhance.Contrast(image)
+    image = enhancer.enhance(contrast)
+    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="Configuration Options"):
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Self-portrait oil painting, a beautiful cyborg with golden hair, 8k")
+            steps_input = gr.Slider(minimum=1, maximum=10, label="Inference Steps", value=4)
+            guidance_input = gr.Slider(minimum=0, maximum=2, label="Guidance Scale", value=1)
+            generate_button = gr.Button("Generate Image")
+    with gr.Row():
+        with gr.Accordion(label="Papercut 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")
+
+    with gr.Row():
+        with gr.Accordion(label="Inpainting"):
+            inpaint_prompt_input = gr.Textbox(label="Prompt for Inpainting", placeholder="a castle on top of a mountain, highly detailed, 8k")
+            init_image_input = gr.File(label="Initial Image")
+            mask_image_input = gr.File(label="Mask Image")
+            inpaint_steps_input = gr.Slider(minimum=1, maximum=10, label="Inference Steps", value=4)
+            inpaint_guidance_input = gr.Slider(minimum=0, maximum=2, label="Guidance Scale", value=1)
+            inpaint_button = gr.Button("Inpaint Image")
+
+    with gr.Row():
+        with gr.Accordion(label="Image Modification (Experimental)"):
+            brightness_slider = gr.Slider(minimum=0.5, maximum=1.5, step=1, label="Brightness")
+            contrast_slider = gr.Slider(minimum=0.5, maximum=1.5, step=1, label="Contrast")
+            modify_button = gr.Button("Modify Image")
+
+    
+
+    generate_button.click(
+        generate_image,
+        inputs=[prompt_input, steps_input, guidance_input],
+        outputs=image_output
+    )
+
+    modify_button.click(
+        modify_image,
+        inputs=[image_output, brightness_slider, contrast_slider],
+        outputs=image_output
+    )
+    inpaint_button.click(
+        inpaint_image,
+        inputs=[inpaint_prompt_input, init_image_input, mask_image_input, inpaint_steps_input, inpaint_guidance_input],
+        outputs=image_output
+    )
+    adapter_generate_button.click(
+        generate_image_with_adapter,
+        inputs=[adapter_prompt_input, adapter_steps_input, adapter_guidance_input],
+        outputs=image_output
+    )
+
 demo.launch()

requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/huggingface/diffusers.git
 git+https://github.com/huggingface/transformers.git
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch
-uuid
+peft
 pydantic
 Pillow
 accelerate
 spaces
-invisible_watermark
+invisible_watermark