pyh-129 commited on
Commit
a44e7f8
·
1 Parent(s): 052b8d9
Files changed (1) hide show
  1. app.py +159 -0
app.py ADDED
@@ -0,0 +1,159 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ import numpy as np
4
+ from diffusers import AutoPipelineForText2Image
5
+ import torch
6
+ import os
7
+ # class SaveImageEveryNStepsCallback:
8
+ # def __init__(self, output_dir, total_steps, interval):
9
+ # self.output_dir = output_dir
10
+ # self.interval = interval
11
+ # if not os.path.exists(output_dir):
12
+ # os.makedirs(output_dir)
13
+ # self.step_list = set(range(0, total_steps, interval))
14
+
15
+ # def __call__(self, scheduler, **kwargs):
16
+ # current_step = kwargs["step"]
17
+ # if current_step in self.step_list:
18
+ # image = kwargs["sample"].detach().cpu().squeeze().permute(1, 2, 0)
19
+ # image = (image + 1) / 2 # normalize image
20
+ # image = image.clamp(0, 1) * 255 # scale to 0-255
21
+ # image = image.numpy().astype("uint8")
22
+ # image_path = os.path.join(self.output_dir, f"image_at_step_{current_step}.png")
23
+ # Image.fromarray(image).save(image_path)
24
+ # print(f"Image saved at step {current_step}")
25
+ output_dir = "./saved_images"
26
+ # def save_image_callback(pipeline, i, t, latents, **kwargs):
27
+ # interval = 5 # Save an image every 5 steps
28
+ # if i % interval == 0:
29
+ # # Convert latents to image
30
+ # image = pipeline.decode_latents_to_image(latents) # Adjust method call according to actual API
31
+ # image = (image + 1) / 2 * 255
32
+ # image = image.clip(0, 255).astype(np.uint8)
33
+ # image = Image.fromarray(image)
34
+ # # Save the image
35
+ # image_path = os.path.join(output_dir, f"image_at_step_{i}.png")
36
+ # image.save(image_path)
37
+ # print(f"Image saved at step {i}")
38
+
39
+ def latents_to_rgb(latents):
40
+ weights = (
41
+ (60, -60, 25, -70),
42
+ (60, -5, 15, -50),
43
+ (60, 10, -5, -35),
44
+ )
45
+ weights_tensor = torch.t(torch.tensor(weights, dtype=latents.dtype).to(latents.device))
46
+ biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(latents.device)
47
+ rgb_tensor = torch.einsum("...lxy,lr -> ...rxy", latents, weights_tensor) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
48
+ image_array = rgb_tensor.clamp(0, 255).byte().cpu().numpy().transpose(1, 2, 0)
49
+ return Image.fromarray(image_array)
50
+
51
+ # Callback function to save images at specific intervals
52
+ def decode_tensors(pipe, step, timestep, callback_kwargs):
53
+ latents = callback_kwargs["latents"]
54
+ image = latents_to_rgb(latents[0])
55
+ image.save(f"./output_images/{step}.png")
56
+ return callback_kwargs
57
+ # 加载预训练模型和权重
58
+ pipeline = AutoPipelineForText2Image.from_pretrained("bguisard/stable-diffusion-nano-2-1", torch_dtype=torch.float16).to("cuda")
59
+ pipeline.load_lora_weights("/root/autodl-tmp/Proj/city_demo/checkpoint-15000",weight_name="pytorch_lora_weights.safetensors")
60
+
61
+ def generate_image(text,option):
62
+ num_steps = 50
63
+ interval = num_steps // 10
64
+ output_dir = "./intermediate_images"
65
+ # callback = SaveImageEveryNStepsCallback(output_dir, num_steps, interval)
66
+ # generator = torch.manual_seed(42)
67
+ # image = pipeline(text, num_inference_steps=num_steps, generator=generator, callback_on_step_end=decode_tensors,
68
+ # callback_on_step_end_tensor_inputs=["latents"])
69
+ while True:
70
+ image = pipeline(text, num_inference_steps=num_steps)
71
+ final_image = image.images[0]
72
+ if option == "Ratio < 5":
73
+ if calculate_building_ratio(final_image) < 5:
74
+ final_pil_image = final_image.convert('L')
75
+ return final_pil_image
76
+ else:
77
+ if calculate_building_ratio(final_image) >= 5:
78
+ final_pil_image = final_image.convert('L')
79
+ return final_pil_image
80
+
81
+ # final_pil_image = Image.fromarray((final_image.cpu().numpy() * 255).astype('uint8'))
82
+
83
+
84
+
85
+ # Save as JPEG
86
+ # image_path = os.path.join(output_dir, "final_image.jpg")
87
+ # final_pil_image.save(image_path, "JPEG")
88
+
89
+ return final_pil_image
90
+ # return final_image
91
+
92
+ # def generate_image(text):
93
+ # # 直接指定图片路径
94
+ # image_path = '/root/autodl-tmp/Proj/city_diffusion_demo/images/beijing/beijing_0.png'
95
+
96
+ # # 加载图片
97
+ # image = Image.open(image_path)
98
+
99
+ # return image
100
+
101
+ def check_requirements(image, requirement):
102
+ # 根据选中的要求检查图片
103
+ # 示例中的需求检查逻辑需要根据具体需求实现
104
+ if requirement == "Option 1":
105
+ # 检查条件1
106
+ pass
107
+ elif requirement == "Option 2":
108
+ # 检查条件2
109
+ pass
110
+ return True # 假设总是返回True
111
+
112
+ def generate_compliant_image(text, requirements):
113
+ while True:
114
+ image = generate_image(text)
115
+ if check_requirements(image, requirements):
116
+ break
117
+ return image
118
+ def calculate_building_ratio(image):
119
+ # 加载图片并转换为灰度图
120
+ # img = Image.open(image_path).convert('L')
121
+ img_array = np.array(image.convert('L'))
122
+
123
+ # 建筑区域定义为所有非零像素
124
+ building_area = np.count_nonzero(img_array != 255)
125
+
126
+ # # 找到最高和最低的有建筑的像素行,用于估算楼层高度
127
+ # non_zero_rows = np.nonzero(img_array)[0]
128
+ # if non_zero_rows.size == 0:
129
+ # return 0 # 如果没有建筑,则返回0
130
+ # min_row, max_row = np.min(non_zero_rows), np.max(non_zero_rows)
131
+ # height = max_row - min_row + 1
132
+ height = np.sum(img_array[img_array != 255])
133
+ print(height)
134
+ print(img_array[img_array != 255])
135
+ # 估算楼层数,假设每层楼高3米
136
+ floors = height / 3
137
+
138
+ # 地块的总面积即整个图像的面积
139
+ total_area = img_array.size
140
+
141
+ # 计算比例:建筑的底面积 * 楼层 / 地块的总面积
142
+ ratio = (floors) / total_area
143
+ print(ratio)
144
+ return ratio/10
145
+
146
+
147
+
148
+ iface = gr.Interface(
149
+ fn=generate_image,
150
+ inputs=[
151
+ gr.Textbox(label="Prompt"),
152
+ gr.Dropdown(choices=["Ratio < 5", "Ratio >= 5"], label="Select Ratio Requirement")
153
+ ],
154
+ outputs="image",
155
+ title="Image of Buildings Generation",
156
+ description="Enter text and specify requirements for the generated image. The image will be regenerated until it meets the requirements."
157
+ )
158
+
159
+ iface.launch(share=True)