Create app.py

app.py

@@ -0,0 +1,47 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from PIL import Image
+
+def load_model():
+    controlnet = ControlNetModel.from_pretrained(
+        "Yuanshi/OminiControl",
+        torch_dtype=torch.float16,
+        use_safetensors=True
+    )
+    
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        controlnet=controlnet,
+        torch_dtype=torch.float16,
+        safety_checker=None
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
+    return pipe
+
+def generate(image, prompt, resolution):
+    pipe = load_model()
+    output = pipe(
+        prompt=prompt,
+        image=image,
+        num_inference_steps=20,
+        controlnet_conditioning_scale=1.0,
+        width=resolution,
+        height=resolution
+    ).images[0]
+    return output
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Enter your prompt"),
+        gr.Radio(choices=[512, 1024], value=512, label="Resolution")
+    ],
+    outputs=gr.Image(label="Generated Image"),
+    title="OminiControl Image Editor",
+    description="Upload an image and provide a prompt to edit it."
+)
+
+if __name__ == "__main__":
+    demo.launch()