Spaces:
Running
on
Zero
Running
on
Zero
bug fix
Browse files
app.py
CHANGED
|
@@ -2,7 +2,7 @@ import spaces
|
|
| 2 |
import gradio as gr
|
| 3 |
import numpy as np
|
| 4 |
import random
|
| 5 |
-
import
|
| 6 |
import torch
|
| 7 |
from diffusers import (
|
| 8 |
ControlNetModel,
|
|
@@ -58,13 +58,14 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
|
|
| 58 |
|
| 59 |
|
| 60 |
def get_depth_map(image):
|
|
|
|
| 61 |
image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
|
| 62 |
with torch.no_grad(), torch.autocast("cuda"):
|
| 63 |
depth_map = depth_estimator(image).predicted_depth
|
| 64 |
-
|
| 65 |
depth_map = torch.nn.functional.interpolate(
|
| 66 |
depth_map.unsqueeze(1),
|
| 67 |
-
size=
|
| 68 |
mode="bicubic",
|
| 69 |
align_corners=False,
|
| 70 |
)
|
|
@@ -79,18 +80,20 @@ def get_depth_map(image):
|
|
| 79 |
|
| 80 |
|
| 81 |
@spaces.GPU(enable_queue=True)
|
| 82 |
-
def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed):
|
| 83 |
|
| 84 |
if image_url:
|
| 85 |
orginal_image = load_image(image_url)
|
| 86 |
else:
|
| 87 |
-
orginal_image =
|
| 88 |
|
| 89 |
size = (orginal_image.size[0], orginal_image.size[1])
|
|
|
|
| 90 |
depth_image = get_depth_map(orginal_image)
|
| 91 |
generator = torch.Generator().manual_seed(seed)
|
| 92 |
generated_image = pipe(
|
| 93 |
prompt=prompt,
|
|
|
|
| 94 |
negative_prompt=n_prompt,
|
| 95 |
width=size[0],
|
| 96 |
height=size[1],
|
|
@@ -98,7 +101,7 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
|
|
| 98 |
num_inference_steps=num_steps,
|
| 99 |
strength=control_strength,
|
| 100 |
generator=generator,
|
| 101 |
-
|
| 102 |
).images[0]
|
| 103 |
return [[depth_image, generated_image], "ok"]
|
| 104 |
|
|
|
|
| 2 |
import gradio as gr
|
| 3 |
import numpy as np
|
| 4 |
import random
|
| 5 |
+
from PIL import Image
|
| 6 |
import torch
|
| 7 |
from diffusers import (
|
| 8 |
ControlNetModel,
|
|
|
|
| 58 |
|
| 59 |
|
| 60 |
def get_depth_map(image):
|
| 61 |
+
original_size = (image.size[1], image.size[0])
|
| 62 |
image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
|
| 63 |
with torch.no_grad(), torch.autocast("cuda"):
|
| 64 |
depth_map = depth_estimator(image).predicted_depth
|
| 65 |
+
print("get_depth_map", original_size)
|
| 66 |
depth_map = torch.nn.functional.interpolate(
|
| 67 |
depth_map.unsqueeze(1),
|
| 68 |
+
size=original_size,
|
| 69 |
mode="bicubic",
|
| 70 |
align_corners=False,
|
| 71 |
)
|
|
|
|
| 80 |
|
| 81 |
|
| 82 |
@spaces.GPU(enable_queue=True)
|
| 83 |
+
def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed, progress=gr.Progress()):
|
| 84 |
|
| 85 |
if image_url:
|
| 86 |
orginal_image = load_image(image_url)
|
| 87 |
else:
|
| 88 |
+
orginal_image = Image.fromarray(image)
|
| 89 |
|
| 90 |
size = (orginal_image.size[0], orginal_image.size[1])
|
| 91 |
+
print(size)
|
| 92 |
depth_image = get_depth_map(orginal_image)
|
| 93 |
generator = torch.Generator().manual_seed(seed)
|
| 94 |
generated_image = pipe(
|
| 95 |
prompt=prompt,
|
| 96 |
+
image=orginal_image,
|
| 97 |
negative_prompt=n_prompt,
|
| 98 |
width=size[0],
|
| 99 |
height=size[1],
|
|
|
|
| 101 |
num_inference_steps=num_steps,
|
| 102 |
strength=control_strength,
|
| 103 |
generator=generator,
|
| 104 |
+
control_image=depth_image
|
| 105 |
).images[0]
|
| 106 |
return [[depth_image, generated_image], "ok"]
|
| 107 |
|