Spaces:
Sleeping
Sleeping
| 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() | |