Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,524 Bytes
5376cf0 847481c 600e217 847481c 0ab8f18 847481c 5376cf0 600e217 0ab8f18 847481c 0ab8f18 600e217 847481c 0ab8f18 847481c 64ab7c4 847481c 64ab7c4 847481c 600e217 af997b5 600e217 af997b5 600e217 af997b5 847481c 600e217 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
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
import io
import os
# Load the base & refiner pipelines
base = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
)
base.to("cuda:0")
# Load your model
pipe = StableDiffusionXLPipeline.from_pretrained(
"segmind/SSD-1B",
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16"
)
pipe.to("cuda:0")
refiner = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0",
text_encoder_2=base.text_encoder_2,
vae=base.vae,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
refiner.to("cuda:0")
refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
@spaces.GPU # Apply the GPU decorator
def generate_and_save_image(prompt, negative_prompt=''):
# Generate image using the provided prompts
image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
# Generate a unique UUID for the filename
unique_id = str(uuid.uuid4())
image_path = f"generated_images/{unique_id}.jpeg"
# Save generated image locally
os.makedirs('generated_images', exist_ok=True)
image.save(image_path, format='JPEG')
# Return the path of the saved image to display in Gradio interface
return image_path
def generate_image_with_refinement(prompt):
n_steps = 40
high_noise_frac = 0.8
# run both experts
image = base(
prompt=prompt,
num_inference_steps=n_steps,
denoising_end=high_noise_frac,
output_type="latent",
).images
image = refiner(
prompt=prompt,
num_inference_steps=n_steps,
denoising_start=high_noise_frac,
image=image,
).images[0]
# Save the image as before
unique_id = str(uuid.uuid4())
image_path = f"generated_images_refined/{unique_id}.jpeg"
os.makedirs('generated_images_refined', exist_ok=True)
image.save(image_path, format='JPEG')
return image_path
# 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.")
with gr.Row():
prompt1 = gr.Textbox(label="Enter prompt")
negative_prompt = gr.Textbox(label="Enter negative prompt (optional)")
generate_button1 = gr.Button("Generate Image")
output_image1 = gr.Image(type="filepath", label="Generated Image")
generate_button1.click(
generate_and_save_image,
inputs=[prompt1, negative_prompt],
outputs=output_image1
)
with gr.Column():
gr.Markdown("## Refined Image Generation")
gr.Markdown("Enter a prompt to generate a refined image.")
with gr.Row():
prompt2 = gr.Textbox(label="Enter prompt for refined generation")
generate_button2 = gr.Button("Generate Refined Image")
output_image2 = gr.Image(type="filepath", label="Generated Refined Image")
generate_button2.click(
generate_image_with_refinement,
inputs=[prompt2],
outputs=output_image2
)
# Launch the combined Gradio app
demo.launch()
|