Upload app.py

app.py

@@ -0,0 +1,58 @@
+import os
+import torch
+from PIL import Image
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+import gradio as gr
+
+# Disable oneDNN custom operations
+os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
+
+# Clear PyTorch cache
+torch.cuda.empty_cache()
+
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if device == "cuda":
+    print("CUDA is available. Device count:", torch.cuda.device_count())
+    print("Current device:", torch.cuda.current_device())
+    print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
+else:
+    print("CUDA is not available. Using CPU.")
+
+# Load ControlNet model with OpenPose pre-trained weights from Hugging Face
+controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16)
+
+# Load the Stable Diffusion model
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
+).to(device)
+
+# Function for inference
+def generate_image(prompt, target_image, pose_image):
+    try:
+        # Resize images
+        target_image = target_image.resize((512, 512))
+        pose_image = pose_image.resize((512, 512))
+
+        # Generate image with ControlNet
+        output = pipe(prompt=prompt, image=target_image, control_image=pose_image, num_inference_steps=50)
+
+        # Return the result
+        return output["sample"][0]
+    except Exception as e:
+        print(f"Error during image generation: {e}")
+        return None
+
+# Setup Gradio Interface
+interface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Image(label="Target Image", type="pil"),
+        gr.Image(label="Pose Image (Reference)", type="pil")
+    ],
+    outputs=gr.Image(label="Generated Image")
+)
+
+# Launch the interface
+interface.launch()