Skip to content

Commit 1fdea05

Browse files
wbrunawbruna
committed
update stable-diffusion.cpp to master-52a97b3
1 parent d70214f commit 1fdea05

32 files changed

+769094
-2045
lines changed

‎otherarch/sdcpp/avi_writer.h‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
#ifndef __AVI_WRITER_H__
2+
#define __AVI_WRITER_H__
3+
4+
#include <stdint.h>
5+
#include <stdio.h>
6+
#include <stdlib.h>
7+
#include <string.h>
8+
9+
#include "stable-diffusion.h"
10+
11+
#ifndef INCLUDE_STB_IMAGE_WRITE_H
12+
#include "stb_image_write.h"
13+
#endif
14+
15+
typedef struct {
16+
uint32_t offset;
17+
uint32_t size;
18+
} avi_index_entry;
19+
20+
// Write 32-bit little-endian integer
21+
void write_u32_le(FILE* f, uint32_t val) {
22+
fwrite(&val, 4, 1, f);
23+
}
24+
25+
// Write 16-bit little-endian integer
26+
void write_u16_le(FILE* f, uint16_t val) {
27+
fwrite(&val, 2, 1, f);
28+
}
29+
30+
/**
31+
* Create an MJPG AVI file from an array of sd_image_t images.
32+
* Images are encoded to JPEG using stb_image_write.
33+
*
34+
* @param filename Output AVI file name.
35+
* @param images Array of input images.
36+
* @param num_images Number of images in the array.
37+
* @param fps Frames per second for the video.
38+
* @param quality JPEG quality (0-100).
39+
* @return 0 on success, -1 on failure.
40+
*/
41+
int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality = 90) {
42+
if (num_images == 0) {
43+
fprintf(stderr, "Error: Image array is empty.\n");
44+
return -1;
45+
}
46+
47+
FILE* f = fopen(filename, "wb");
48+
if (!f) {
49+
perror("Error opening file for writing");
50+
return -1;
51+
}
52+
53+
uint32_t width = images[0].width;
54+
uint32_t height = images[0].height;
55+
uint32_t channels = images[0].channel;
56+
if (channels != 3 && channels != 4) {
57+
fprintf(stderr, "Error: Unsupported channel count: %u\n", channels);
58+
fclose(f);
59+
return -1;
60+
}
61+
62+
// --- RIFF AVI Header ---
63+
fwrite("RIFF", 4, 1, f);
64+
long riff_size_pos = ftell(f);
65+
write_u32_le(f, 0); // Placeholder for file size
66+
fwrite("AVI ", 4, 1, f);
67+
68+
// 'hdrl' LIST (header list)
69+
fwrite("LIST", 4, 1, f);
70+
write_u32_le(f, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
71+
fwrite("hdrl", 4, 1, f);
72+
73+
// 'avih' chunk (AVI main header)
74+
fwrite("avih", 4, 1, f);
75+
write_u32_le(f, 56);
76+
write_u32_le(f, 1000000 / fps); // Microseconds per frame
77+
write_u32_le(f, 0); // Max bytes per second
78+
write_u32_le(f, 0); // Padding granularity
79+
write_u32_le(f, 0x110); // Flags (HASINDEX | ISINTERLEAVED)
80+
write_u32_le(f, num_images); // Total frames
81+
write_u32_le(f, 0); // Initial frames
82+
write_u32_le(f, 1); // Number of streams
83+
write_u32_le(f, width * height * 3); // Suggested buffer size
84+
write_u32_le(f, width);
85+
write_u32_le(f, height);
86+
write_u32_le(f, 0); // Reserved
87+
write_u32_le(f, 0); // Reserved
88+
write_u32_le(f, 0); // Reserved
89+
write_u32_le(f, 0); // Reserved
90+
91+
// 'strl' LIST (stream list)
92+
fwrite("LIST", 4, 1, f);
93+
write_u32_le(f, 4 + 8 + 56 + 8 + 40);
94+
fwrite("strl", 4, 1, f);
95+
96+
// 'strh' chunk (stream header)
97+
fwrite("strh", 4, 1, f);
98+
write_u32_le(f, 56);
99+
fwrite("vids", 4, 1, f); // Stream type: video
100+
fwrite("MJPG", 4, 1, f); // Codec: Motion JPEG
101+
write_u32_le(f, 0); // Flags
102+
write_u16_le(f, 0); // Priority
103+
write_u16_le(f, 0); // Language
104+
write_u32_le(f, 0); // Initial frames
105+
write_u32_le(f, 1); // Scale
106+
write_u32_le(f, fps); // Rate
107+
write_u32_le(f, 0); // Start
108+
write_u32_le(f, num_images); // Length
109+
write_u32_le(f, width * height * 3); // Suggested buffer size
110+
write_u32_le(f, (uint32_t)-1); // Quality
111+
write_u32_le(f, 0); // Sample size
112+
write_u16_le(f, 0); // rcFrame.left
113+
write_u16_le(f, 0); // rcFrame.top
114+
write_u16_le(f, 0); // rcFrame.right
115+
write_u16_le(f, 0); // rcFrame.bottom
116+
117+
// 'strf' chunk (stream format: BITMAPINFOHEADER)
118+
fwrite("strf", 4, 1, f);
119+
write_u32_le(f, 40);
120+
write_u32_le(f, 40); // biSize
121+
write_u32_le(f, width);
122+
write_u32_le(f, height);
123+
write_u16_le(f, 1); // biPlanes
124+
write_u16_le(f, 24); // biBitCount
125+
fwrite("MJPG", 4, 1, f); // biCompression (FOURCC)
126+
write_u32_le(f, width * height * 3); // biSizeImage
127+
write_u32_le(f, 0); // XPelsPerMeter
128+
write_u32_le(f, 0); // YPelsPerMeter
129+
write_u32_le(f, 0); // Colors used
130+
write_u32_le(f, 0); // Colors important
131+
132+
// 'movi' LIST (video frames)
133+
long movi_list_pos = ftell(f);
134+
fwrite("LIST", 4, 1, f);
135+
long movi_size_pos = ftell(f);
136+
write_u32_le(f, 0); // Placeholder for movi size
137+
fwrite("movi", 4, 1, f);
138+
139+
avi_index_entry* index = (avi_index_entry*)malloc(sizeof(avi_index_entry) * num_images);
140+
if (!index) {
141+
fclose(f);
142+
return -1;
143+
}
144+
145+
// Encode and write each frame as JPEG
146+
struct {
147+
uint8_t* buf;
148+
size_t size;
149+
} jpeg_data;
150+
151+
for (int i = 0; i < num_images; i++) {
152+
jpeg_data.buf = NULL;
153+
jpeg_data.size = 0;
154+
155+
// Callback function to collect JPEG data into memory
156+
auto write_to_buf = [](void* context, void* data, int size) {
157+
auto jd = (decltype(jpeg_data)*)context;
158+
jd->buf = (uint8_t*)realloc(jd->buf, jd->size + size);
159+
memcpy(jd->buf + jd->size, data, size);
160+
jd->size += size;
161+
};
162+
163+
// Encode to JPEG in memory
164+
stbi_write_jpg_to_func(
165+
write_to_buf,
166+
&jpeg_data,
167+
images[i].width,
168+
images[i].height,
169+
channels,
170+
images[i].data,
171+
quality);
172+
173+
// Write '00dc' chunk (video frame)
174+
fwrite("00dc", 4, 1, f);
175+
write_u32_le(f, jpeg_data.size);
176+
index[i].offset = ftell(f) - 8;
177+
index[i].size = jpeg_data.size;
178+
fwrite(jpeg_data.buf, 1, jpeg_data.size, f);
179+
180+
// Align to even byte size
181+
if (jpeg_data.size % 2)
182+
fputc(0, f);
183+
184+
free(jpeg_data.buf);
185+
}
186+
187+
// Finalize 'movi' size
188+
long cur_pos = ftell(f);
189+
long movi_size = cur_pos - movi_size_pos - 4;
190+
fseek(f, movi_size_pos, SEEK_SET);
191+
write_u32_le(f, movi_size);
192+
fseek(f, cur_pos, SEEK_SET);
193+
194+
// Write 'idx1' index
195+
fwrite("idx1", 4, 1, f);
196+
write_u32_le(f, num_images * 16);
197+
for (int i = 0; i < num_images; i++) {
198+
fwrite("00dc", 4, 1, f);
199+
write_u32_le(f, 0x10);
200+
write_u32_le(f, index[i].offset);
201+
write_u32_le(f, index[i].size);
202+
}
203+
204+
// Finalize RIFF size
205+
cur_pos = ftell(f);
206+
long file_size = cur_pos - riff_size_pos - 4;
207+
fseek(f, riff_size_pos, SEEK_SET);
208+
write_u32_le(f, file_size);
209+
fseek(f, cur_pos, SEEK_SET);
210+
211+
fclose(f);
212+
free(index);
213+
214+
return 0;
215+
}
216+
217+
#endif // __AVI_WRITER_H__

0 commit comments

Comments
 (0)