Deadmon commited on
Commit
5ae0057
·
verified ·
1 Parent(s): cdc4413

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +86 -0
app.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import gradio as gr
4
+ import numpy as np
5
+ from PIL import Image
6
+ from einops import rearrange
7
+ import requests
8
+ from diffusers.utils import load_image
9
+ from diffusers import FluxControlNetPipeline, FluxControlNetModel
10
+ from gradio_imageslider import ImageSlider
11
+
12
+ # Pretrained model paths
13
+ base_model = 'black-forest-labs/FLUX.1-dev'
14
+ controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
15
+
16
+ # Load the ControlNet and pipeline models
17
+ controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
18
+ pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
19
+ pipe.to("cuda")
20
+
21
+ # Define control modes
22
+ CONTROL_MODES = {
23
+ 0: "Canny",
24
+ 1: "Tile",
25
+ 2: "Depth",
26
+ 3: "Blur",
27
+ 4: "Pose",
28
+ 5: "Gray (Low)",
29
+ 6: "LQ"
30
+ }
31
+
32
+ def preprocess_image(image, target_width, target_height):
33
+ image = image.resize((target_width, target_height), Image.LANCZOS)
34
+ return image
35
+
36
+ def generate_image(prompt, control_image, control_mode, controlnet_conditioning_scale, num_steps, guidance, width, height, seed, random_seed):
37
+ if random_seed:
38
+ seed = np.random.randint(0, 10000)
39
+
40
+ # Ensure width and height are multiples of 16
41
+ width = 16 * (width // 16)
42
+ height = 16 * (height // 16)
43
+
44
+ # Set the seed for reproducibility
45
+ torch.manual_seed(seed)
46
+
47
+ # Preprocess control image
48
+ control_image = preprocess_image(control_image, width, height)
49
+
50
+ # Generate the image with the selected control mode and other parameters
51
+ with torch.no_grad():
52
+ image = pipe(
53
+ prompt,
54
+ control_image=control_image,
55
+ control_mode=control_mode,
56
+ width=width,
57
+ height=height,
58
+ controlnet_conditioning_scale=controlnet_conditioning_scale,
59
+ num_inference_steps=num_steps,
60
+ guidance_scale=guidance
61
+ ).images[0]
62
+
63
+ return image
64
+
65
+ # Define the Gradio interface
66
+ interface = gr.Interface(
67
+ fn=generate_image,
68
+ inputs=[
69
+ gr.Textbox(label="Prompt"),
70
+ gr.Image(type="pil", label="Control Image"),
71
+ gr.Dropdown(choices=[f"{i}: {name}" for i, name in CONTROL_MODES.items()], type="index", label="Control Mode", default=0),
72
+ gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="ControlNet Conditioning Scale"),
73
+ gr.Slider(step=1, minimum=1, maximum=64, value=24, label="Num Steps"),
74
+ gr.Slider(minimum=0.1, maximum=10, value=3.5, label="Guidance"),
75
+ gr.Slider(minimum=128, maximum=1024, step=128, value=512, label="Width"),
76
+ gr.Slider(minimum=128, maximum=1024, step=128, value=512, label="Height"),
77
+ gr.Number(value=42, label="Seed"),
78
+ gr.Checkbox(label="Random Seed")
79
+ ],
80
+ outputs=ImageSlider(label="Generated Image"),
81
+ title="FLUX.1 Controlnet with Multiple Modes",
82
+ description="Generate images using ControlNet and a text prompt with adjustable control modes."
83
+ )
84
+
85
+ if __name__ == "__main__":
86
+ interface.launch()