Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,404 Bytes
0a82683 1ec1764 0a82683 1ec1764 0a82683 1ec1764 0a82683 1ec1764 0a82683 |
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 |
import gradio as gr
import torch
from PIL import Image
from src.pipeline_pe_clone import FluxPipeline
import spaces
# Load default image from assets as an example
default_image = Image.open("assets/1.png")
@spaces.GPU()
def generate_image(model_path, image, height, width, prompt, guidance_scale, num_steps, lora_name):
# Load the model
pipeline = FluxPipeline.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
).to('cuda')
# Load and fuse base LoRA weights
pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name="pretrain.safetensors")
pipeline.fuse_lora()
pipeline.unload_lora_weights()
# Load selected LoRA effect if not using the pretrained base model
if lora_name != 'pretrained':
pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name=f"{lora_name}.safetensors")
# Prepare the input image
condition_image = image.resize((height, width)).convert("RGB")
# Generate the output image
result = pipeline(
prompt=prompt,
condition_image=condition_image,
height=height,
width=width,
guidance_scale=guidance_scale,
num_inference_steps=num_steps,
max_sequence_length=512
).images[0]
return result
# Create Gradio interface with sliders for numeric inputs
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Model Path", value="black-forest-labs/FLUX.1-dev"),
gr.Image(label="Input Image", type="pil", value=default_image),
gr.Slider(label="Height", value=768, minimum=256, maximum=1024, step=64),
gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=64),
gr.Textbox(label="Prompt", value="add a halo and wings for the cat by sksmagiceffects"),
gr.Slider(label="Guidance Scale", value=3.5, minimum=1.0, maximum=10.0, step=0.1),
gr.Slider(label="Number of Steps", value=20, minimum=1, maximum=100, step=1),
gr.Dropdown(
label="LoRA Name",
choices=["pretrained", "sksmagiceffects", "sksmonstercalledlulu",
"skspaintingeffects", "sksedgeeffect", "skscatooneffect"],
value="sksmagiceffects"
)
],
outputs=gr.Image(label="Output Image", type="pil"),
title="FLUX Image Generation with LoRA"
)
if __name__ == "__main__":
iface.launch()
|