Add video cancellation feature and update API integration

- Implemented a new endpoint to cancel ongoing video generation tasks in `video_generation.py`.
- Enhanced the frontend to manage task IDs and cancellation state in `GenerationContext.tsx` and `GenerationForm.tsx`.
- Updated the UI to allow users to cancel video generation with a button in `GenerationProgress.tsx`.
- Added support for handling cancellation in the video generation process.
- Included new dependencies in `requirements.txt` for OpenAI Whisper.

api/image_service.py

@@ -26,7 +26,7 @@ async def upload_image(file: UploadFile = File(...)):
         data_url = f"data:{file.content_type};base64,{encoded}"
         
         # Get public URL from env or use default
-        public_url = os.getenv('PUBLIC_URL', 'http://localhost:4000')
+        public_url = os.getenv('VITE_API_BASE_URL', 'http://localhost:4000')
         
         # Compress and store, get hosted URL
         hosted_url = await compress_and_store_image(data_url, public_url)

api/video_generation.py

@@ -15,6 +15,7 @@ from datetime import datetime
 
 from utils.image_processor import compress_and_store_image
 from utils.storage import video_results, sse_clients, cleanup_old_results
+import asyncio
 
 router = APIRouter()
 
@@ -575,6 +576,44 @@ async def get_video_status(task_id: str):
             detail=f"Failed to check video status: {error_msg}"
         )
 
+@router.post("/veo/cancel/{task_id}")
+async def cancel_video_generation(task_id: str):
+    """
+    Cancel an ongoing video generation task
+    """
+    try:
+        # Send cancellation event to SSE clients
+        if task_id in sse_clients:
+            queue = sse_clients[task_id]
+            await send_sse_event(task_id, {
+                'status': 'cancelled',
+                'message': 'Video generation cancelled by user'
+            })
+            # Close the SSE connection
+            del sse_clients[task_id]
+            print(f"✅ Cancelled video generation: {task_id}")
+        
+        # Mark result as cancelled
+        if task_id in video_results:
+            video_results[task_id] = {
+                'code': 499,  # Client Closed Request
+                'msg': 'Video generation cancelled by user',
+                'taskId': task_id,
+                'timestamp': datetime.now().timestamp()
+            }
+        
+        return JSONResponse(
+            status_code=200,
+            content={'code': 200, 'msg': 'Video generation cancelled', 'taskId': task_id}
+        )
+    
+    except Exception as e:
+        print(f"❌ Error cancelling video generation: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to cancel video generation: {str(e)}"
+        )
+
 @router.get("/veo/download")
 async def download_video(url: str):
     """

frontend/src/components/GenerationForm.tsx

@@ -37,8 +37,8 @@ const aspectRatios = ['9:16', '16:9', '1:1'];
 type GenerationMode = 'extend' | 'frame-continuity';
 
 export const GenerationForm: React.FC<GenerationFormProps> = ({ provider, onBack }) => {
-  const { startGeneration, updateProgress, addVideo, setStep, setError, setRetryState, updateSegments, state } = useGeneration();
-  const { retryState, generatedVideos, segments } = state;
+  const { startGeneration, updateProgress, addVideo, setStep, setError, setRetryState, updateSegments, addTaskId, removeTaskId, state } = useGeneration();
+  const { retryState, generatedVideos, segments, isCancelling } = state;
   
   // Draft storage key
   const draftKey = `video-gen-draft-${provider}`;
@@ -554,6 +554,11 @@ export const GenerationForm: React.FC<GenerationFormProps> = ({ provider, onBack
 
         // Generate video with automatic retry (retries once on failure)
         updateProgress(`Processing video ${i + 1}... (this may take 1-2 minutes)`);
+        // Check if cancelled
+        if (isCancelling) {
+          throw new Error('Generation cancelled by user');
+        }
+
         const videoUrl = await generateVideoWithRetry(async () => {
         if (i === 0 || (i === startIndex && startIndex > 0)) {
           // First segment OR resuming after failure: use generate API with current image
@@ -567,6 +572,7 @@ export const GenerationForm: React.FC<GenerationFormProps> = ({ provider, onBack
             voiceType: formState.voiceType,
           });
           currentTaskId = generateResult.taskId;
+          addTaskId(currentTaskId);
             return generateResult;
         } else {
           // Subsequent segments: use extend API
@@ -577,6 +583,7 @@ export const GenerationForm: React.FC<GenerationFormProps> = ({ provider, onBack
               formState.voiceType
             );
           currentTaskId = extendResult.taskId;
+          addTaskId(currentTaskId);
             return extendResult;
           }
         }, 300000, (attempt) => {
@@ -602,6 +609,11 @@ export const GenerationForm: React.FC<GenerationFormProps> = ({ provider, onBack
         });
 
         updateProgress(`Completed video ${i + 1} of ${payload.segments.length}`, i + 1, payload.segments.length);
+        
+        // Remove task ID after completion
+        if (currentTaskId) {
+          removeTaskId(currentTaskId);
+        }
       }
 
       clearDraft(); // Clear draft on successful generation
@@ -612,15 +624,23 @@ export const GenerationForm: React.FC<GenerationFormProps> = ({ provider, onBack
       console.error('Generation error:', err);
       const errorMessage = err instanceof Error ? err.message : 'Generation failed';
       
-      // Enable retry mode
-      setRetryState({
-        failedSegmentIndex: generatedVideos.length, // Current segment that failed
-        error: errorMessage
-      });
-      setStep('configuring'); // Go back to form, but with retry overlay
+      // If cancelled, don't show retry option
+      if (errorMessage.includes('cancelled') || isCancelling) {
+        setError('Generation cancelled by user');
+        setStep('error');
+      } else {
+        // Enable retry mode
+        setRetryState({
+          failedSegmentIndex: generatedVideos.length, // Current segment that failed
+          error: errorMessage
+        });
+        setStep('configuring'); // Go back to form, but with retry overlay
+      }
       
     } finally {
       setIsGenerating(false);
+      // Clean up any remaining task IDs
+      state.activeTaskIds.forEach(taskId => removeTaskId(taskId));
     }
   };
 

frontend/src/components/GenerationProgress.tsx

@@ -61,9 +61,15 @@ interface ActivityLog {
   icon?: 'video' | 'audio' | 'brain' | 'download' | 'image';
 }
 
+const XIcon = () => (
+  <svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+    <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
+  </svg>
+);
+
 export const GenerationProgress: React.FC = () => {
-  const { state } = useGeneration();
-  const { progress, provider, generatedVideos, segments } = state;
+  const { state, cancelGeneration } = useGeneration();
+  const { progress, provider, generatedVideos, segments, isCancelling, activeTaskIds } = state;
   
   const [elapsedTime, setElapsedTime] = useState(0);
   const [activityLog, setActivityLog] = useState<ActivityLog[]>([]);
@@ -176,18 +182,37 @@ export const GenerationProgress: React.FC = () => {
           <div>
             <h2 className="text-2xl font-bold text-void-100">Generating Videos</h2>
             <p className="text-void-400 text-sm mt-1">
-              {provider === 'kling' ? 'Kling AI' : 'Replicate'} • {segments.length} segments
+              {provider === 'kling' ? 'KIE API' : 'Replicate'} • {segments.length} segments
             </p>
           </div>
-          <div className="flex items-center gap-6 text-sm">
-            <div className="flex items-center gap-2 text-void-400">
-              <ClockIcon />
-              <span>Elapsed: <span className="text-void-200 font-mono">{formatTime(elapsedTime)}</span></span>
-            </div>
-            {progress.current > 0 && remainingTime > 0 && (
+          <div className="flex items-center gap-4">
+            <div className="flex items-center gap-6 text-sm">
               <div className="flex items-center gap-2 text-void-400">
-                <span>Est. remaining: <span className={`font-mono ${accentClass}`}>{formatTime(remainingTime)}</span></span>
+                <ClockIcon />
+                <span>Elapsed: <span className="text-void-200 font-mono">{formatTime(elapsedTime)}</span></span>
               </div>
+              {progress.current > 0 && remainingTime > 0 && (
+                <div className="flex items-center gap-2 text-void-400">
+                  <span>Est. remaining: <span className={`font-mono ${accentClass}`}>{formatTime(remainingTime)}</span></span>
+                </div>
+              )}
+            </div>
+            {activeTaskIds.length > 0 && (
+              <button
+                onClick={cancelGeneration}
+                disabled={isCancelling}
+                className={`
+                  px-4 py-2 rounded-lg font-medium text-sm transition-all
+                  flex items-center gap-2
+                  ${isCancelling 
+                    ? 'bg-void-700 text-void-400 cursor-not-allowed' 
+                    : 'bg-red-500/20 hover:bg-red-500/30 text-red-400 hover:text-red-300 border border-red-500/30'
+                  }
+                `}
+              >
+                <XIcon />
+                {isCancelling ? 'Cancelling...' : 'Cancel Generation'}
+              </button>
             )}
           </div>
         </div>

frontend/src/context/GenerationContext.tsx

@@ -22,6 +22,8 @@ const initialState: GenerationState = {
   error: null,
   taskId: null,
   retryState: null,
+  activeTaskIds: [] as string[],
+  isCancelling: false,
 };
 
 // Action types
@@ -35,6 +37,9 @@ type GenerationAction =
   | { type: 'SET_ERROR'; payload: string | null }
   | { type: 'SET_TASK_ID'; payload: string | null }
   | { type: 'SET_RETRY_STATE'; payload: { failedSegmentIndex: number; error: string } | null }
+  | { type: 'ADD_TASK_ID'; payload: string }
+  | { type: 'REMOVE_TASK_ID'; payload: string }
+  | { type: 'SET_CANCELLING'; payload: boolean }
   | { type: 'RESET' };
 
 // Reducer
@@ -64,6 +69,12 @@ function generationReducer(state: GenerationState, action: GenerationAction): Ge
       return { ...state, taskId: action.payload };
     case 'SET_RETRY_STATE':
       return { ...state, retryState: action.payload };
+    case 'ADD_TASK_ID':
+      return { ...state, activeTaskIds: [...state.activeTaskIds, action.payload] };
+    case 'REMOVE_TASK_ID':
+      return { ...state, activeTaskIds: state.activeTaskIds.filter(id => id !== action.payload) };
+    case 'SET_CANCELLING':
+      return { ...state, isCancelling: action.payload };
     case 'RESET':
       return { ...initialState, provider: state.provider };
     default:
@@ -86,6 +97,9 @@ interface GenerationContextValue {
   setError: (error: string | null) => void;
   setRetryState: (state: { failedSegmentIndex: number; error: string } | null) => void;
   updateSegments: (segments: VeoSegment[]) => void;
+  addTaskId: (taskId: string) => void;
+  removeTaskId: (taskId: string) => void;
+  cancelGeneration: () => Promise<void>;
   reset: () => void;
 }
 
@@ -148,6 +162,41 @@ export function GenerationProvider({ children }: { children: ReactNode }) {
       dispatch({ type: 'SET_SEGMENTS', payload: segments });
     },
 
+    addTaskId: (taskId) => {
+      dispatch({ type: 'ADD_TASK_ID', payload: taskId });
+    },
+
+    removeTaskId: (taskId) => {
+      dispatch({ type: 'REMOVE_TASK_ID', payload: taskId });
+    },
+
+    cancelGeneration: async () => {
+      dispatch({ type: 'SET_CANCELLING', payload: true });
+      try {
+        const { klingCancel } = await import('@/utils/api');
+        // Cancel all active tasks
+        const currentTaskIds = [...state.activeTaskIds];
+        const cancelPromises = currentTaskIds.map(taskId => 
+          klingCancel(taskId).catch(err => {
+            console.warn(`Failed to cancel task ${taskId}:`, err);
+          })
+        );
+        await Promise.all(cancelPromises);
+        // Clear all task IDs
+        currentTaskIds.forEach(id => {
+          dispatch({ type: 'REMOVE_TASK_ID', payload: id });
+        });
+        dispatch({ type: 'SET_TASK_ID', payload: null });
+        dispatch({ type: 'SET_ERROR', payload: 'Generation cancelled by user' });
+        dispatch({ type: 'SET_STEP', payload: 'error' });
+      } catch (error) {
+        console.error('Error cancelling generation:', error);
+        dispatch({ type: 'SET_ERROR', payload: 'Failed to cancel generation' });
+      } finally {
+        dispatch({ type: 'SET_CANCELLING', payload: false });
+      }
+    },
+
     reset: () => {
       dispatch({ type: 'RESET' });
     },

frontend/src/types/index.ts

@@ -183,6 +183,8 @@ export interface GenerationState {
     failedSegmentIndex: number;
     error: string;
   } | null;
+  activeTaskIds: string[];
+  isCancelling: boolean;
 }
 
 export interface GeneratedVideo {

frontend/src/utils/api.ts

@@ -210,6 +210,13 @@ export async function klingGetStatus(taskId: string): Promise<VideoStatusRespons
   return apiRequest<VideoStatusResponse>(`/api/veo/status/${taskId}`);
 }
 
+export async function klingCancel(taskId: string): Promise<{ code: number; msg: string; taskId: string }> {
+  return apiRequest<{ code: number; msg: string; taskId: string }>(`/api/veo/cancel/${taskId}`, {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+  });
+}
+
 export function createKlingEventSource(taskId: string): EventSource {
   return new EventSource(`${API_BASE}/api/veo/events/${taskId}`);
 }
@@ -501,6 +508,10 @@ export function waitForKlingVideo(taskId: string, timeoutMs: number = 300000): P
           clearTimeout(timeout);
           eventSource.close();
           resolve(data.url);
+        } else if (data.status === 'cancelled') {
+          clearTimeout(timeout);
+          eventSource.close();
+          reject(new Error('Video generation cancelled by user'));
         } else if (data.status === 'failed' || data.code !== undefined && data.code !== 200) {
           clearTimeout(timeout);
           eventSource.close();

requirements.txt

@@ -36,3 +36,4 @@ python-multipart==0.0.17
 # pytest==8.3.0
 # pytest-asyncio==0.24.0
 
+openai-whisper