changess

app.py

@@ -28,11 +28,7 @@ with gr.Blocks(theme=theme) as GenDemo:
         with gr.Row(variant="panel"):
             with gr.Column():
                 prompt = gr.Textbox(label="Enter a discription of a shoe")
-                negative_prompt = gr.Textbox(label="Negative Prompt", value="low quality, bad quality, sketches, legs")
-                select = gr.Dropdown(label="Select a model", choices=["Canny","Depth","Normal"])
-            
-                scale = gr.Slider(label="Control Image Scale", minimum=0.1, maximum=1.0, step=0.1, value=0.5, visible=(select == "Canny"))
-                
+                select = gr.Dropdown(label="Select a model", choices=["Canny","Depth","Normal"])      
                 controlNet_image = gr.Image(label="Enter an image of a shoe, that you want to use as a reference", type='numpy')
                 gr.Examples(
                        examples=[
@@ -46,8 +42,7 @@ with gr.Blocks(theme=theme) as GenDemo:
                 button_gen = gr.Button("Generate Image", elem_id="generateIm", variant="primary")
                 gen_image = gr.Image(label="Generated Image", image_mode="RGBA", type='pil', show_download_button=True, show_label=False)
         
-        select.change()
-        button_gen.click(check_prompt, inputs=[prompt]).success(generate_image, inputs=[prompt, negative_prompt, controlNet_image, scale], outputs=[gen_image])
+        button_gen.click(check_prompt, inputs=[prompt]).success(generate_image, inputs=[prompt, negative_prompt, controlNet_image, select], outputs=[gen_image])
 
     with gr.Tab("Image to 3D Model Generator"):
         with gr.Row(variant="panel"):

util/text_img.py

@@ -1,8 +1,9 @@
 import spaces
 import rembg
 import torch
-from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, AutoencoderKL
 import cv2
+from transformers import pipeline
 import numpy as np
 from PIL import Image
 import gradio as gr
@@ -14,47 +15,94 @@ def check_prompt(prompt):
     if prompt is None:
         raise gr.Error("Please enter a prompt!")
 
-controlnet = ControlNetModel.from_pretrained(
-    "diffusers/controlnet-canny-sdxl-1.0",
-    torch_dtype=torch.float16,
-    use_safetensors=True
-)
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, use_safetensors=True)
-pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    controlnet=controlnet,
-    vae=vae,
-    torch_dtype=torch.float16,
-    use_safetensors=True
-)
+controlNet_normal = ControlNetModel.from_pretrained(
+        "fusing/stable-diffusion-v1-5-controlnet-normal", 
+        torch_dtype=torch.float16
+    )
+
+controlNet_depth = ControlNetModel.from_pretrained(
+        "lllyasviel/sd-controlnet-depth", 
+        torch_dtype=torch.float16
+    )
+controlNet_MAP = {"Normal": controlNet_normal, "Depth": controlNet_depth}
 
-pipe.to("cuda")
+# vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, use_safetensors=True)
 
 # Function to generate an image from text using diffusion
 @spaces.GPU
-def generate_image(prompt, negative_prompt, control_image, scale=0.5):
+def generate_image(prompt, control_image, controlnet):
     prompt += "no background, side view, minimalist shot, single shoe, no legs, product photo"
     
-    canny_image = get_canny(control_image)
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlNet_MAP[controlnet],
+    torch_dtype=torch.float16,
+    safety_checker = None
+    )
+    
+    pipe.to("cuda")
 
-    image = pipe(
-        prompt,
+    if controlnet == "Normal":
+        control_image = get_normal(control_image)
+    elif controlnet == "Depth":
+        control_image = get_depth(control_image)
+    
+    image = pipe(prompt,
         negative_prompt=negative_prompt,
-        image=canny_image,
-        controlnet_conditioning_scale=scale,
-    ).images[0]
+        image=control_image,
+        controlnet_conditioning_scale=1.0).images[0]
+
     image2 = rembg.remove(image)
 
     return image2
 
-def get_canny(image):
-    image = np.array(image)
 
-    low_threshold = 100
-    high_threshold = 200
+def get_normal(image):
+    depth_estimator = pipeline("depth-estimation", model ="Intel/dpt-hybrid-midas" )
+
+    image = depth_estimator(image)['predicted_depth'][0]
+
+    image = image.numpy()
+
+    image_depth = image.copy()
+    image_depth -= np.min(image_depth)
+    image_depth /= np.max(image_depth)
 
-    image = cv2.Canny(image,low_threshold,high_threshold)
-    image = image[:,:,None]
+    bg_threhold = 0.4
+
+    x = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
+    x[image_depth < bg_threhold] = 0
+
+    y = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)
+    y[image_depth < bg_threhold] = 0
+
+    z = np.ones_like(x) * np.pi * 2.0
+
+    image = np.stack([x, y, z], axis=2)
+    image /= np.sum(image ** 2.0, axis=2, keepdims=True) ** 0.5
+    image = (image * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
+    normalimage = Image.fromarray(image)
+
+    return normalimage
+
+def get_depth(image):
+    depth_estimator = pipeline('depth-estimation')
+
+    image = depth_estimator(image)['depth']
+    image = np.array(image)
+    image = image[:, :, None]
     image = np.concatenate([image, image, image], axis=2)
-    canny_image = Image.fromarray(image)
-    return canny_image
+    depthimage = Image.fromarray(image)
+    return depthimage
+
+# def get_canny(image):
+#     image = np.array(image)
+
+#     low_threshold = 100
+#     high_threshold = 200
+
+#     image = cv2.Canny(image,low_threshold,high_threshold)
+#     image = image[:,:,None]
+#     image = np.concatenate([image, image, image], axis=2)
+#     canny_image = Image.fromarray(image)
+#     return canny_image