Create app.py

app.py

@@ -0,0 +1,63 @@
+import torch
+import numpy as np
+import gradio as gr
+from transformers import pipeline
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers.utils import load_image, make_image_grid
+from PIL import Image
+
+# Function to get depth map
+def get_depth_map(image, depth_estimator):
+    image = depth_estimator(image)["depth"]
+    image = np.array(image)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    detected_map = torch.from_numpy(image).float() / 255.0
+    depth_map = detected_map.permute(2, 0, 1)
+    return depth_map
+
+# Main function to process the image and prompt
+def process_image_and_prompt(input_image, prompt):
+    # Convert PIL Image to the format expected by the pipeline
+    input_image = input_image.convert("RGB")
+    
+    # Load depth estimator
+    depth_estimator = pipeline("depth-estimation")
+
+    # Get depth map
+    depth_map = get_depth_map(input_image, depth_estimator).unsqueeze(0).half().to("cuda")
+    
+    # Load the ControlNet model and the StableDiffusionControlNetImg2ImgPipeline
+    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-normal", torch_dtype=torch.float16, use_safetensors=True)
+    pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        controlnet=controlnet,
+        torch_dtype=torch.float16,
+        use_safetensors=True
+    ).to("cuda")
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_model_cpu_offload()
+    
+    # Generate the image
+    output = pipe(
+        prompt,
+        image=input_image,
+        control_image=depth_map,
+    ).images[0]
+    
+    # Convert output to PIL Image for Gradio display
+    output_image = Image.fromarray(output)
+    
+    return input_image, output_image
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=process_image_and_prompt,
+    inputs=[gr.inputs.Image(type="pil"), gr.inputs.Textbox(label="Prompt")],
+    outputs=[gr.outputs.Image(label="Original Image"), gr.outputs.Image(label="Generated Image")],
+    title="Image and Prompt Processing with Stable Diffusion",
+    description="Upload an image and enter a prompt to generate a new image."
+)
+
+# Launch the Gradio app
+iface.launch()