jbilcke-hf HF staff commited on
Commit
458a627
1 Parent(s): 3637e81

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +70 -6
handler.py CHANGED
@@ -69,6 +69,70 @@ apply_dirty_hack_to_patch_file_extensions_and_bypass_filter("/repository")
69
  #logger.info("💡 Printing directory structure of ""/repository"":")
70
  #print_directory_structure("/repository")
71
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
  @dataclass
73
  class GenerationConfig:
74
  """Configuration for video generation"""
@@ -339,12 +403,12 @@ class EndpointHandler:
339
 
340
  # Check if image-to-video generation is requested
341
  if input_image:
342
- # Process base64 image
343
- if input_image.startswith('data:'):
344
- input_image = input_image.split(',', 1)[1]
345
- image_bytes = base64.b64decode(input_image)
346
- image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
347
- generation_kwargs["image"] = image
348
  frames = self.image_to_video(**generation_kwargs).frames
349
  else:
350
  frames = self.text_to_video(**generation_kwargs).frames
 
69
  #logger.info("💡 Printing directory structure of ""/repository"":")
70
  #print_directory_structure("/repository")
71
 
72
+
73
+ def process_input_image(image_data: str, target_width: int, target_height: int) -> Image.Image:
74
+ """
75
+ Process input image from base64, resize and crop to target dimensions
76
+
77
+ Args:
78
+ image_data: Base64 encoded image data
79
+ target_width: Desired width
80
+ target_height: Desired height
81
+
82
+ Returns:
83
+ Processed PIL Image
84
+ """
85
+ try:
86
+ # Handle data URI format
87
+ if image_data.startswith('data:'):
88
+ image_data = image_data.split(',', 1)[1]
89
+
90
+ # Decode base64
91
+ image_bytes = base64.b64decode(image_data)
92
+ image = Image.open(io.BytesIO(image_bytes))
93
+
94
+ # Convert to RGB if necessary
95
+ if image.mode not in ('RGB', 'RGBA'):
96
+ image = image.convert('RGB')
97
+ elif image.mode == 'RGBA':
98
+ # Handle transparency by compositing on white background
99
+ background = Image.new('RGB', image.size, (255, 255, 255))
100
+ background.paste(image, mask=image.split()[3])
101
+ image = background
102
+
103
+ # Calculate target aspect ratio
104
+ target_aspect = target_width / target_height
105
+
106
+ # Get current dimensions
107
+ orig_width, orig_height = image.size
108
+ orig_aspect = orig_width / orig_height
109
+
110
+ # Calculate dimensions for resizing
111
+ if orig_aspect > target_aspect:
112
+ # Image is wider than target
113
+ new_height = target_height
114
+ new_width = int(target_height * orig_aspect)
115
+ else:
116
+ # Image is taller than target
117
+ new_width = target_width
118
+ new_height = int(target_width / orig_aspect)
119
+
120
+ # Resize image
121
+ image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
122
+
123
+ # Center crop to target dimensions
124
+ left = (new_width - target_width) // 2
125
+ top = (new_height - target_height) // 2
126
+ right = left + target_width
127
+ bottom = top + target_height
128
+
129
+ image = image.crop((left, top, right, bottom))
130
+
131
+ return image
132
+
133
+ except Exception as e:
134
+ raise ValueError(f"Failed to process input image: {str(e)}")
135
+
136
  @dataclass
137
  class GenerationConfig:
138
  """Configuration for video generation"""
 
403
 
404
  # Check if image-to-video generation is requested
405
  if input_image:
406
+ processed_image = process_input_image(
407
+ input_image,
408
+ config.width,
409
+ config.height
410
+ )
411
+ generation_kwargs["image"] = processed_image
412
  frames = self.image_to_video(**generation_kwargs).frames
413
  else:
414
  frames = self.text_to_video(**generation_kwargs).frames