File size: 1,859 Bytes
2b9dd70
 
271f923
2b9dd70
 
 
 
 
 
 
4f62c4b
 
fb3a978
 
bab5d39
 
4f62c4b
 
 
bab5d39
 
 
271f923
b9dbdf6
bab5d39
 
 
 
 
 
 
4f62c4b
 
 
2b9dd70
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from diffusers import FluxPipeline
import torch
import spaces

@spaces.GPU
def generate_image(prompt, num_inference_steps, guidance_scale):
    # Load the base model
    base_model_id = "black-forest-labs/FLUX.1-dev"
    #pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16)

    #model_path = "MegaTronX/SuicideGirl-FLUX"
    model_path = "https://huggingface.co/MegaTronX/SuicideGirl-FLUX/blob/main/SuicideGirls.safetensors"
    
    '''pipe = StableDiffusionPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float16,
    use_safetensors=True
    '''

    # Load the base model and apply LoRA weights using from_single_file
    pipe = FluxPipeline.from_single_file(
        "https://huggingface.co/MegaTronX/SuicideGirl-FLUX/blob/main/SuicideGirls.safetensors",
        torch_dtype=torch.float16,
        use_safetensors=True,
        load_safety_checker=False,
        variant="fp16",
        original_config_file=base_model_id
    )


        # Load the LoRA weights
    #pipe.unet.load_attn_procs("MegaTronX/SuicideGirl-FLUX")
    
    pipe = pipe.to("cuda")
    
    # Generate the image
    image = pipe(
        prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale
    ).images[0]
    
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(minimum=1, maximum=100, value=50, label="Number of Inference Steps"),
        gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Image Generation with Custom LoRA",
    description="Generate images using a custom LoRA model trained on Flux Dev."
)

iface.launch()