Spaces:
Sleeping
Sleeping
File size: 2,336 Bytes
1bff21e |
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 |
import gradio as gr
from diffusers import StableDiffusionPipeline
def apply_lora(pipeline, lora_path):
"""
Dummy function to simulate the application of LoRA weights.
Replace this with your actual code to load and integrate LoRA weights.
"""
if lora_path:
print(f"Applying LoRA weights from {lora_path}")
# Insert your LoRA integration code here.
return pipeline
def generate_image(model_name, lora_path, width, height, inference_steps, prompt):
# Use the provided model name or fall back to a default model
model_id = model_name.strip() if model_name.strip() else "CompVis/stable-diffusion-v1-5"
# Load the diffusion pipeline from Hugging Face
pipeline = StableDiffusionPipeline.from_pretrained(model_id)
device = "cuda" if gr.get_config().get("device") == "gpu" else "cpu"
pipeline = pipeline.to(device)
# Apply LoRA if a path is provided
if lora_path.strip():
pipeline = apply_lora(pipeline, lora_path.strip())
# Generate the image using the specified parameters
result = pipeline(prompt, width=width, height=height, num_inference_steps=inference_steps)
return result.images[0]
# Build the Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Image Generator with Custom Model & LoRA")
model_name_box = gr.Textbox(
label="Enter Model Name/ID (e.g., CompVis/stable-diffusion-v1-5)",
value="CompVis/stable-diffusion-v1-5",
lines=1
)
lora_path_box = gr.Textbox(
label="Enter LoRA Path (leave empty if not using)",
value="",
lines=1
)
width_slider = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Image Width")
height_slider = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Image Height")
steps_slider = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps")
prompt_box = gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
generate_button = gr.Button("Generate Image")
output_image = gr.Image(label="Generated Image")
generate_button.click(
fn=generate_image,
inputs=[model_name_box, lora_path_box, width_slider, height_slider, steps_slider, prompt_box],
outputs=output_image
)
demo.launch()
|