File size: 1,571 Bytes
a46020e db79ba5 a46020e 6ef9c62 a46020e 6ef9c62 a46020e db79ba5 a46020e |
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 |
import torch
import random
import gradio as gr
from diffusers import StableDiffusionPipeline, DDIMScheduler
# Fine-tuned model weights path
model_path = "model"
# Load fine-tuned pipeline
pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32)
# Change Scheduler to DDIMScheduler
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
# For memory optimization -- only for GPU
# pipe.enable_xformers_memory_efficient_attention()
# Generate the image from the given text prompt
def generate_response(prompt):
# define the negative prompt
negative_prompt = "bad anatomy, ugly, deformed, desfigured, distorted face, poorly drawn, blurry, low quality, low definition, lowres, out of frame, out of image, cropped, cut off, signature, watermark"
# initialize the variables
num_samples = 1
guidance_scale = 7.5
num_inference_steps = 10
height = 512
width = 512
seed = random.randint(0, 2147483647)
generator = torch.Generator(device='cpu').manual_seed(seed)
with torch.inference_mode():
imgs = pipe(
prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_images_per_prompt=num_samples,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator
).images
return imgs[0]
# Create and launch the Gradio UI
gradio_ui = gr.Interface(fn=generate_response, inputs="text", outputs="image")
gradio_ui.launch()
|