diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs
index cfe9b466..1a822ec2 100644
--- a/libwayshot/src/dispatch.rs
+++ b/libwayshot/src/dispatch.rs
@@ -213,12 +213,6 @@ pub struct CaptureFrameState {
     pub(crate) session_done: bool,
 }
 
-impl CaptureFrameState {
-    pub(crate) fn dmabuf_formats(&self) -> &[DMAFrameFormat] {
-        &self.dmabuf_formats
-    }
-}
-
 impl Dispatch<ZwpLinuxDmabufV1, ()> for CaptureFrameState {
     fn event(
         _frame: &mut Self,
@@ -311,12 +305,16 @@ impl Dispatch<ExtImageCopyCaptureSessionV1, ()> for CaptureFrameState {
                 set_format.format = format;
             }
             ext_image_copy_capture_session_v1::Event::DmabufFormat { format, .. } => {
+                let mut width = 0;
+                let mut height = 0;
+                if !state.formats.is_empty() {
+                    let format = state.formats[0];
+                    width = format.size.width;
+                    height = format.size.height;
+                }
                 state.dmabuf_formats.push(DMAFrameFormat {
                     format,
-                    size: Size {
-                        width: 0,
-                        height: 0,
-                    },
+                    size: Size { width, height },
                 });
             }
             ext_image_copy_capture_session_v1::Event::Done => {
diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs
index 7ffe638b..72032b25 100644
--- a/libwayshot/src/lib.rs
+++ b/libwayshot/src/lib.rs
@@ -595,22 +595,15 @@ impl WayshotConnection {
         target: &WayshotTarget,
         capture_region: Option<EmbeddedRegion>,
     ) -> Result<(DMAFrameFormat, DMAFrameGuard, BufferObject<()>)> {
-        let WayshotTarget::Screen(output) = target else {
-            return Err(Error::Unsupported(
-                "ext-image-copy support is on road, so here we restrict it on WlOutput".to_owned(),
-            ));
-        };
         match &self.dmabuf_state {
             Some(dmabuf_state) => {
-                // FIXME: now it won't work on ext-image-copy
-                let (state, event_queue, frame) = self
-                    .capture_output_frame_get_state_wlr_dmabuf_tmp(
-                        cursor_overlay as i32,
-                        output,
-                        capture_region,
-                    )?;
-                //self.capture_target_frame_get_state(cursor_overlay, target, capture_region)?;
-                let frame_format = state.dmabuf_formats()[0];
+                let (state, event_queue, frame) =
+                    self.capture_target_frame_get_state(cursor_overlay, target, capture_region)?;
+                if state.dmabuf_formats.is_empty() {
+                    return Err(Error::NoSupportedBufferFormat);
+                }
+
+                let frame_format = state.dmabuf_formats[0];
                 tracing::trace!("Selected frame buffer format: {:#?}", frame_format);
                 let gbm = &dmabuf_state.gbmdev;
                 let bo = gbm.create_buffer_object::<()>(
@@ -731,7 +724,7 @@ impl WayshotConnection {
             .unwrap_or(Options::PaintCursors);
         let session = manager.create_session(&source, options, &qh, ());
         let frame = session.create_frame(&qh, ());
-        while state.formats.is_empty() {
+        while !state.session_done {
             event_queue.blocking_dispatch(&mut state)?;
         }
         tracing::trace!(
@@ -801,36 +794,6 @@ impl WayshotConnection {
         }
     }
 
-    // NOTE: because in 0.4.0, the dmabuf cannot work on ext-image-copy
-    // So here force it go to the wlr-screencopy
-    // Will be deleted in the feature
-    fn capture_output_frame_get_state_wlr_dmabuf_tmp(
-        &self,
-        cursor_overlay: i32,
-        output: &WlOutput,
-        capture_region: Option<EmbeddedRegion>,
-    ) -> Result<(
-        CaptureFrameState,
-        EventQueue<CaptureFrameState>,
-        WayshotFrame,
-    )> {
-        let state = CaptureFrameState {
-            formats: Vec::new(),
-            dmabuf_formats: Vec::new(),
-            state: None,
-            buffer_done: AtomicBool::new(false),
-            toplevels: Vec::new(),
-            session_done: false,
-        };
-        let event_queue = self.conn.new_event_queue::<CaptureFrameState>();
-        self.capture_output_frame_get_state_wlr(
-            state,
-            event_queue,
-            cursor_overlay,
-            output,
-            capture_region,
-        )
-    }
     // This API is exposed to provide users with access to window manager (WM)
     // information. For instance, enabling Vulkan in wlroots alters the display
     // format. Consequently, using PipeWire to capture streams without knowing
@@ -1608,7 +1571,7 @@ impl WayshotConnection {
         };
         let session = manager.create_session(&source, options, &qh, ());
         let frame = session.create_frame(&qh, ());
-        while state.formats.is_empty() {
+        while !state.session_done {
             event_queue.blocking_dispatch(&mut state)?;
         }