Upload 2 files

app.py

@@ -0,0 +1,46 @@
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForImageTextToText
+import torch
+
+model_id = "zai-org/UI2Code_N"
+
+processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForImageTextToText.from_pretrained(
+    model_id, torch_dtype=torch.bfloat16, device_map="auto"
+)
+
+def generate_code(image, prompt):
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": prompt},
+        ]
+    }]
+
+    inputs = processor.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        return_dict=True, return_tensors="pt"
+    ).to(model.device)
+
+    generated_ids = model.generate(**inputs, max_new_tokens=4096)
+    output = processor.decode(
+        generated_ids[0][inputs["input_ids"].shape[1]:],
+        skip_special_tokens=True
+    )
+    return output
+
+demo = gr.Interface(
+    fn=generate_code,
+    inputs=[
+        gr.Image(type="filepath", label="Upload UI Image"),
+        gr.Textbox(label="Prompt", value="Generate HTML for this UI")
+    ],
+    outputs=gr.Code(label="Generated Code", language="html"),
+    title="UI2Codeⁿ Demo",
+    description="A Visual Language Model for Interactive UI-to-Code Generation."
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+

requirements.txt

@@ -0,0 +1,2 @@
+gradio
+