File size: 1,512 Bytes
6455ecf
 
f5bd834
6455ecf
f5bd834
 
 
 
 
 
 
 
 
 
 
 
 
6455ecf
f5bd834
 
 
 
 
 
 
 
6455ecf
f5bd834
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6455ecf
f5bd834
 
6455ecf
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
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

def generate_image(prompt, guidance_scale=7.5, num_steps=30, lora_scale=0.8):
    try:
        # Initialize the pipeline
        pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
        )
        
        # Load LoRA weights
        pipe.load_lora_weights("Scalino84/my-flux-face-v2")
        
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
        
        # Generate image
        image = pipe(
            prompt=prompt,
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
        ).images[0]
        
        return image
    
    except Exception as e:
        print(f"Error: {str(e)}")
        raise gr.Error(str(e))

# Create Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="a photo of xyz person, professional headshot"),
        gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale"),
        gr.Slider(minimum=20, maximum=100, value=30, label="Number of Steps"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.8, label="LoRA Scale")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Flux Face Generator",
    description="Generate images using your custom LoRA model"
)

# Launch the interface
interface.launch()