Update app.py

app.py

@@ -3,22 +3,31 @@ from PIL import Image
 import torch
 from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
 
+# Load model from HF Hub (your model repo)
+model_id = "yutengz/ip2p-RoboPredict"
+
 pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
-    "yutengz/ip2p-RoboPredict",
+    model_id,
     torch_dtype=torch.float16,
-    safety_checker=None
+    safety_checker=None,
 ).to("cuda")
+
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-def generate(prompt, input_image):
-    edited = pipe(prompt=prompt, image=input_image.convert("RGB").resize((256, 256))).images[0]
-    return edited
+def predict(image: Image.Image, prompt: str):
+    image = image.convert("RGB").resize((256, 256))
+    result = pipe(image=image, prompt=prompt).images[0]
+    return result
 
-gr.Interface(
-    fn=generate,
-    inputs=[gr.Text(label="Instruction Prompt"), gr.Image(label="Input Image")],
-    outputs=gr.Image(label="Edited Image"),
-    title="InstructPix2Pix Finetuned Model",
-    description="Upload an image and enter an instruction to see how the model transforms it."
-).launch()
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="Source Image"),
+        gr.Textbox(label="Instruction Prompt", placeholder="e.g., stack the blocks"),
+    ],
+    outputs=gr.Image(label="Predicted Image"),
+    title="🧠 ip2p-RoboPredict",
+    description="A fine-tuned InstructPix2Pix model for robotic action frame prediction."
+)
 
+demo.launch()