pyh-129
commited on
Commit
·
a44e7f8
1
Parent(s):
052b8d9
add
Browse files
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)
|