File size: 28,344 Bytes
36dd645
9a4edab
1e9882c
9a4edab
 
1e9882c
 
9a4edab
 
524f3f2
75efe1b
9a4edab
d9bc1e6
346f789
 
f9bfec2
524f3f2
d9bc1e6
524f3f2
1e9882c
 
 
 
d9bc1e6
 
8352ac4
 
 
524f3f2
9a4edab
1e9882c
 
c9dba17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1298b7e
c9dba17
 
 
4fbcde7
 
 
 
c9dba17
 
 
 
 
1e9882c
 
a2398fe
1e9882c
 
 
 
 
 
9a4edab
1e9882c
 
2e40bb4
346f789
1e9882c
 
 
 
 
 
 
 
 
 
640b59e
 
 
1e9882c
 
 
 
 
 
 
 
 
 
 
 
9a4edab
1e9882c
 
 
b45d081
1e9882c
01fbb1a
 
 
c9dba17
1e9882c
 
 
 
 
 
60f7312
1e9882c
 
 
 
13fad22
1e9882c
 
bcf5c7a
1e9882c
 
 
 
 
2bc3ef7
f326d9e
 
67e82c5
2bc3ef7
1e9882c
 
 
2e40bb4
 
2bc3ef7
cb81547
 
2e40bb4
cb81547
1e9882c
d9bc1e6
 
a84ba92
 
d9bc1e6
 
 
2e40bb4
1e9882c
 
9a4edab
1e9882c
 
 
 
9a4edab
 
 
 
1e9882c
9a4edab
1e9882c
80ca3a0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e2a3dde
 
42dd4f6
e2a3dde
80ca3a0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e2a3dde
42dd4f6
a82a77f
80ca3a0
 
 
 
 
 
 
 
 
 
 
 
 
 
48e039b
80ca3a0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7192cc2
80ca3a0
 
 
 
 
 
 
 
 
 
 
 
191a9de
1e9882c
 
 
 
0b42297
75efe1b
2f73600
75efe1b
e2a3dde
75efe1b
 
db6bd44
2e40bb4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
db6bd44
2e40bb4
 
 
 
 
 
 
 
346f789
 
 
 
 
 
 
 
 
 
 
 
2e40bb4
1298b7e
2105d3a
346f789
 
2e40bb4
 
 
 
346f789
 
 
 
2e40bb4
 
346f789
2e40bb4
 
346f789
 
 
 
80ca3a0
 
 
346f789
2e40bb4
0e95549
 
d9bc1e6
 
01fbb1a
 
 
 
 
 
 
 
0e95549
6bb96e8
0e95549
01fbb1a
 
0e95549
 
 
 
 
 
 
 
 
 
 
6bb96e8
 
01fbb1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1298b7e
01fbb1a
 
 
 
 
 
 
 
c5f9bd0
0e95549
6bb96e8
01fbb1a
 
 
 
 
 
 
 
 
 
d6f45c3
01fbb1a
 
 
 
 
74393b1
01fbb1a
74393b1
 
 
01fbb1a
80ca3a0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e95549
 
 
01fbb1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
80ca3a0
01fbb1a
 
 
 
 
 
80ca3a0
01fbb1a
 
 
 
 
 
 
 
 
 
74393b1
 
 
01fbb1a
 
 
 
 
 
 
 
 
 
 
 
 
0e95549
 
01fbb1a
1298b7e
2e40bb4
d9bc1e6
2e40bb4
 
 
d9bc1e6
2e40bb4
 
0e95549
4b736df
0e95549
d9bc1e6
 
0e95549
 
 
 
 
 
 
 
 
 
 
4b736df
6bb96e8
4b736df
2e40bb4
d9bc1e6
2e40bb4
d9bc1e6
2105d3a
346f789
2105d3a
 
 
 
 
346f789
2105d3a
 
346f789
2105d3a
 
 
2e40bb4
2105d3a
 
d9bc1e6
2105d3a
 
562fba7
2105d3a
 
bd972fb
2105d3a
 
 
 
 
 
ab345d8
2105d3a
 
ab345d8
2105d3a
 
 
 
 
346f789
2105d3a
 
c5f9bd0
0e95549
6bb96e8
2105d3a
2e40bb4
2105d3a
 
 
 
346f789
2105d3a
 
2e40bb4
9abb18b
2105d3a
5fa4d82
 
2105d3a
 
 
5fa4d82
 
 
 
2105d3a
 
5fa4d82
 
 
2105d3a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
01fbb1a
2105d3a
 
0e95549
1298b7e
0e95549
 
1298b7e
 
78e3b99
 
 
 
1298b7e
78e3b99
2105d3a
 
1298b7e
2105d3a
 
 
 
 
2e40bb4
1298b7e
 
2105d3a
 
d9bc1e6
2105d3a
 
d468814
 
 
2e40bb4
01fbb1a
1e9882c
01fbb1a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
import os
import io
import json
import base64
import random
import urllib.request
import urllib.parse
import websocket
import uuid
from dotenv import load_dotenv
from flask import Flask, request, jsonify, render_template, send_file, send_from_directory
from PIL import Image
from werkzeug.utils import secure_filename
import urllib.parse
import urllib.request
import time

# Load environment variables from the .env file
load_dotenv()

# Initialize Flask app
app = Flask(__name__)

ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'webp'}  # Define supported image types

# Set server and websocket addresses from environment variables
server_address = os.getenv("SERVER_ADDRESS")
ws_address = os.getenv("WS_ADDRESS")

# Generate a unique client ID
client_id = str(uuid.uuid4())

def allowed_file(filename):
    """Check if the uploaded file has an allowed extension."""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def save_base64_image(b64_string):
    """Decode a base64 string and save it as an image in the static folder."""
    try:
        # Handle Data URI scheme if present
        if ',' in b64_string:
            header, encoded = b64_string.split(',', 1)
            ext = header.split('/')[1].split(';')[0] if '/' in header else 'png'
        else:
            encoded = b64_string
            ext = 'png'

        # Decode the image data
        image_data = base64.b64decode(encoded)

        # Generate a unique path for the image in the static folder
        image_path = f"static/{uuid.uuid4()}.{ext}"

        # Ensure directory exists
        os.makedirs('static', exist_ok=True)

        # Save the image
        with open(image_path, 'wb') as f:
            f.write(image_data)

        print(f"Image saved at: {image_path}", flush=True)

        # Return the path and URL of the saved image
        image_url = f"https://gosign-de-comfyui-api.hf.space/{image_path}"

        print(f"Image path (local): {image_path}", flush=True)
        print(f"Image URL (public): {image_url}", flush=True)

        return image_path, image_url

    except Exception as e:
        raise ValueError(f"Failed to save image: {e}")

def get_image(filename, subfolder, image_type, token):
    url_values = {'filename': filename, 'subfolder': subfolder, 'type': image_type}
    url = f"{server_address}/view?{urllib.parse.urlencode(url_values)}"
    req = urllib.request.Request(url)
    req.add_header("Authorization", f"Bearer {token}")
    try:
        return urllib.request.urlopen(req).read()
    except urllib.error.HTTPError as e:
        print(f"HTTP Error: {e.code} - {e.reason}")
        print(e.read())
        raise

def get_images(ws, workflow, token):
    prompt_id = queue_prompt(workflow, token)
    output_images = {}

    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break  # Execution is done
                    
        # Sleep for 3 seconds before the next iteration
        time.sleep(3)

    history = get_history(prompt_id, token)[prompt_id]
    for node_id in history['outputs']:
        node_output = history['outputs'][node_id]
        images_output = []
        if 'images' in node_output:
            for image in node_output['images']:
                image_data = get_image(image['filename'], image['subfolder'], image['type'], token)
                images_output.append(image_data)
        output_images[node_id] = images_output

    return output_images

# Default route for home welcome
@app.route('/')
def home():
    return render_template('home.html')

                ################################################
                # Generate text to image using FLUX1.DEV Model #
                ################################################

# Generate image route
@app.route('/generate_image', methods=['POST'])
def generate_image():
    data = request.json

    # Extract the token from the request headers
    token = request.headers.get('Authorization')

    if token is None:
        return jsonify({'error': 'No token provided'}), 400
    if token.startswith("Bearer "):
        token = token.split(" ")[1]

    # Base64 decode the encoded token
    # token = base64.b64decode(token).decode("utf-8")

    if 'text_prompt' not in data:
        return jsonify({'error': 'No text prompt provided'}), 400

    text_prompt = data['text_prompt']

    # Get the path to the current file's directory
    current_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(current_dir, 'workflows/flux1_dev_checkpoint_workflow_api.json')

    with open(file_path, 'r', encoding='utf-8') as file:
        workflow_jsondata = file.read()

    workflow = json.loads(workflow_jsondata)
    workflow["6"]["inputs"]["text"] = text_prompt

    # Generate a random 15-digit seed as an integer
    seednum = random.randint(100000000000000, 999999999999999)
    workflow["31"]["inputs"]["seed"] = seednum

    ws = websocket.WebSocket()

    try:
        ws.connect(f"{ws_address}?clientId={client_id}&token={token}", header=
        {"Authorization": f"Bearer {token}"})
    except websocket.WebSocketException as e:
        return jsonify({'error': f'WebSocket connection failed: {str(e)}'}), 500

    images = get_images(ws, workflow, token)
    ws.close()

    output_images_base64 = []

    for node_id in images:
        for image_data in images[node_id]:
            image = Image.open(io.BytesIO(image_data))
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
            output_images_base64.append(img_str)

    return jsonify({'images': output_images_base64})


                ###################################################
                # Edit image with text prompt using OmniGen Model #
                ###################################################

# Route: OmniGen image to image
@app.route('/omnigen/image_to_image', methods=['POST'])
def omnigen_image_to_image():
    data = request.json

    # Extract and validate token
    token = request.headers.get('Authorization')
    if not token or not token.startswith("Bearer "):
        return jsonify({'error': 'Valid Bearer token required'}), 400
    token = token.split(" ")[1]

    # Validate text prompt
    text_prompt = data.get('text_prompt')
    if not text_prompt or not text_prompt.strip():
        return jsonify({'error': 'Text prompt is required'}), 400

    steps = data.get('steps')
    if not steps:
        steps = 50

    image_url = data.get('image_url')
    if not image_url:
        return jsonify({'error': 'image_url is required'}), 400

    # Handle uploaded image or base64 image
    image_file = request.files.get('image')
    base64_image = data.get('base64_image')

    image_path = None  # Initialize image path

    try:
        if image_file:
            # Check if the file has an allowed extension
            if not allowed_file(image_file.filename):
                return jsonify({'error': 'Unsupported image format'}), 400

            # Secure the filename
            filename = secure_filename(image_file.filename)

            # Generate a unique path for the image
            unique_filename = f"{uuid.uuid4()}_{filename}"
            image_path = os.path.join('static', unique_filename)

            # Ensure the 'static' directory exists
            os.makedirs('static', exist_ok=True)

            # Save the image to the static directory
            image_file.save(image_path)

            # Construct the public URL to access the image
            image_url = f"https://gosign-de-comfyui-api.hf.space/{image_path}"

        elif base64_image:
            # Save base64 image
            try:
                pass
                # image_path, image_url = save_base64_image(base64_image)
                # image_url = "https://drive.google.com/uc?id=1JEHEy0zCVWOob4421hLQIPMbO_ebeCPS&export=download"
            except Exception as e:
                raise ValueError(f'Invalid base64 image data: {str(e)}')

        else:
            return jsonify({'error': 'Image is required (file or base64)'}), 400

        # Load workflow configuration
        current_dir = os.path.dirname(os.path.abspath(__file__))
        workflow_path = os.path.join(current_dir, 'workflows/omnigen_image_to_image_workflow_api.json')
        with open(workflow_path, 'r', encoding='utf-8') as f:
            workflow = json.load(f)

        # Modify workflow with inputs
        workflow["6"]["inputs"]["prompt"] = "in image_1 " + text_prompt
        workflow["6"]["inputs"]["num_inference_steps"] = steps
        workflow["12"]["inputs"]["url"] = image_url

        # WebSocket connection to queue the prompt
        ws = websocket.WebSocket()
        ws.connect(f"{ws_address}?clientId={client_id}&token={token}",
                   header={"Authorization": f"Bearer {token}"})

        images = get_images(ws, workflow, token)
        ws.close()

        output_images_base64 = []

        for node_id in images:
            for image_data in images[node_id]:
                image = Image.open(io.BytesIO(image_data))
                buffered = io.BytesIO()
                image.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
                output_images_base64.append(img_str)

        return jsonify({'images': output_images_base64}), 200

    except Exception as e:
        return jsonify({'message': 'Unable to connect to the server. Make sure the server is running', 'error': str(e)}), 500

    finally:
        pass
        # Always delete the image if it was saved
        if image_path and os.path.exists(image_path):
            os.remove(image_path)
            print(f"Deleted temporary image: {image_path}", flush=True)


# Get image route
@app.route('/get_image/<filename>', methods=['GET'])
def get_image_file(filename):
    return send_file(filename, mimetype='image/png')


# Route to serve images
@app.route('/static/<path:filename>', methods=['GET'])
def serve_static(filename):
    print(f"Request for static file: {filename}", flush=True)
    return send_from_directory('static', filename)

# Make a request route
def make_request(url, data=None, headers=None):
    req = urllib.request.Request(url, data=data, headers=headers)
    try:
        with urllib.request.urlopen(req) as response:
            response_body = response.read().decode()  # Decode the response
            # print(response_body)
            return json.loads(response_body)  # Convert to JSON if valid
    except urllib.error.HTTPError as e:
        print(f"HTTPError: {e.code}, {e.reason}")
        print(e.read().decode())  # Print detailed error response
    except urllib.error.URLError as e:
        print(f"URLError: {e.reason}")

# Helper: Queue the prompt
def queue_prompt(workflow, token):
    payload = {"prompt": workflow, "client_id": client_id}
    headers = {
        'Authorization': f'Bearer {token}',
        'Content-Type': 'application/json'
    }
    response = make_request(f"{server_address}/prompt", data=json.dumps(payload).encode('utf-8'), headers=headers)
    if not response or 'prompt_id' not in response:
        raise ValueError("Failed to queue the prompt. Check the request or API response.")
    return response['prompt_id']

# Get ComfyUI prompt history
def get_history(prompt_id, token):
    headers = {
        'Authorization': f'Bearer {token}',
        'Content-Type': 'application/json'
    }
    return make_request(f"{server_address}/history/{prompt_id}", headers=headers)

def get_video_data(filename, subfolder, token):
    """
    Retrieve a video from the server using filename, subfolder, and token.
    """
    # Handle empty subfolder case gracefully
    subfolder = subfolder or ''  # Default to empty string if None

     # Construct query parameters
    url_values = {
        'filename': filename
    }

    # Build the URL with encoded query parameters
    url = f"{server_address}/view?{urllib.parse.urlencode(url_values)}"

    print(f"Requesting URL: {url}", flush=True)

    # Prepare the request with authorization token
    req = urllib.request.Request(url)
    req.add_header("Authorization", f"Bearer {token}")

    try:
        # Fetch and return the video data
        response = urllib.request.urlopen(req)
        return response.read()

    except urllib.error.HTTPError as e:
        print(f"HTTP Error: {e.code} - {e.reason}")
        print(e.read().decode())  # Decode error message for readability
        raise

    except urllib.error.URLError as e:
        print(f"URL Error: {e.reason}")
        raise

                ########################################################
                # Generate image to video using CogVideoX-5B-12V Model #
                ########################################################

# Route: Image to Video
@app.route('/v1/image_to_video', methods=['POST'])
def v1_image_to_video():
    data = request.json

    # Extract and validate token
    token = request.headers.get('Authorization')
    if not token or not token.startswith("Bearer "):
        return jsonify({'error': 'Valid Bearer token required'}), 400
    token = token.split(" ")[1]

    # Validate text prompt
    text_prompt = data.get('text_prompt')
    frame_rate = data.get('frame_rate')
    steps = data.get('steps')
    if not text_prompt or not text_prompt.strip():
        return jsonify({'error': 'Text prompt is required'}), 400

    # Check if frame_rate is missing or invalid
    if not frame_rate:  # If frame_rate is None, empty, or 0
        frame_rate = 24  # Default to 24 fps
    else:
        try:
            frame_rate = int(frame_rate)
            if frame_rate not in [8, 12, 24]:  # Ensure it's one of the allowed values
                return jsonify({'error': 'Frame rate must be a valid number (8, 12, or 24).'}), 400
        except ValueError:
            return jsonify({'error': 'Frame rate must be a valid number (8, 12, or 24).'}), 400

    if not steps:
        steps = 50
    # Handle uploaded image or base64 image
    image_file = request.files.get('image')
    base64_image = data.get('base64_image')

    image_path = None  # Initialize image path

    try:
        if image_file:
            # Check if the file has an allowed extension
            if not allowed_file(image_file.filename):
                return jsonify({'error': 'Unsupported image format'}), 400

            # Secure the filename
            filename = secure_filename(image_file.filename)

            # Generate a unique path for the image
            unique_filename = f"{uuid.uuid4()}_{filename}"
            image_path = os.path.join('static', unique_filename)

            # Ensure the 'static' directory exists
            os.makedirs('static', exist_ok=True)

            # Save the image to the static directory
            image_file.save(image_path)

            # Construct the public URL to access the image
            image_url = f"https://gosign-de-comfyui-api.hf.space/{image_path}"

        elif base64_image:
            # Save base64 image
            try:
                image_path, image_url = save_base64_image(base64_image)
            except Exception as e:
                raise ValueError(f'Invalid base64 image data: {str(e)}')

        else:
            return jsonify({'error': 'Image is required (file or base64)'}), 400

        # Load workflow configuration
        current_dir = os.path.dirname(os.path.abspath(__file__))
        workflow_path = os.path.join(current_dir, 'workflows/cogvideox_image_to_video_workflow_api.json')
        with open(workflow_path, 'r', encoding='utf-8') as f:
            workflow = json.load(f)

        # Modify workflow with inputs
        workflow["30"]["inputs"]["prompt"] = text_prompt
        workflow["31"]["inputs"]["prompt"] = "Low quality, watermark, strange motion, blur"
        workflow["44"]["inputs"]["frame_rate"] = frame_rate
        workflow["57"]["inputs"]["steps"] = steps
        workflow["73"]["inputs"]["url"] = image_url

        # WebSocket connection to queue the prompt
        ws = websocket.WebSocket()
        ws.connect(f"{ws_address}?clientId={client_id}&token={token}",
                   header={"Authorization": f"Bearer {token}"})

        # Queue the prompt
        prompt_id = queue_prompt(workflow, token)

        return jsonify({'prompt_id': prompt_id, 'message': 'Prompt queued successfully', 'get_video_url': f'https://gosign-de-comfyui-api.hf.space/v1/video_tasks/{prompt_id}'}), 202

    except Exception as e:
        return jsonify({'message': 'Unbale to connect to the server. Make sure the server is running', 'error': str(e)}), 500

    finally:
        pass
        # Always delete the image if it was saved
        # if image_path and os.path.exists(image_path):
        #     os.remove(image_path)
        #     print(f"Deleted temporary image: {image_path}", flush=True)


                ###################################################
                # Generate text to video using CogVideoX-5B Model #
                ###################################################

# Route: Text to Video
@app.route('/v1/text_to_video', methods=['POST'])
def v1_text_to_video():
    data = request.json

    # Extract and validate token
    token = request.headers.get('Authorization')
    if not token or not token.startswith("Bearer "):
        return jsonify({'error': 'Valid Bearer token required'}), 400
    token = token.split(" ")[1]

    # Validate text prompt
    text_prompt = data.get('text_prompt')
    frame_rate = data.get('frame_rate')
    steps = data.get('steps')
    if not text_prompt or not text_prompt.strip():
        return jsonify({'error': 'Text prompt is required'}), 400

    # Check if frame_rate is missing or invalid
    if not frame_rate:  # If frame_rate is None, empty, or 0
        frame_rate = 24  # Default to 24 fps
    else:
        try:
            frame_rate = int(frame_rate)
            if frame_rate not in [8, 12, 24]:  # Ensure it's one of the allowed values
                return jsonify({'error': 'Frame rate must be a valid number (8, 12, or 24).'}), 400
        except ValueError:
            return jsonify({'error': 'Frame rate must be a valid number (8, 12, or 24).'}), 400

    if not steps:
        steps = 50

    try:
        # Load workflow configuration
        current_dir = os.path.dirname(os.path.abspath(__file__))
        workflow_path = os.path.join(current_dir, 'workflows/cogvideox_text_to_video_workflow_api.json')
        with open(workflow_path, 'r', encoding='utf-8') as f:
            workflow = json.load(f)

        # Modify workflow with inputs
        workflow["30"]["inputs"]["prompt"] = text_prompt
        workflow["31"]["inputs"]["prompt"] = "Low quality, watermark, strange motion, blur"
        workflow["33"]["inputs"]["frame_rate"] = frame_rate
        workflow["34"]["inputs"]["steps"] = steps

        # WebSocket connection to queue the prompt
        ws = websocket.WebSocket()
        ws.connect(f"{ws_address}?clientId={client_id}&token={token}",
                   header={"Authorization": f"Bearer {token}"})

        # Queue the prompt
        prompt_id = queue_prompt(workflow, token)

        return jsonify({'prompt_id': prompt_id, 'message': 'Prompt queued successfully', 'get_video_url': f'https://gosign-de-comfyui-api.hf.space/v1/video_tasks/{prompt_id}'}), 202

    except Exception as e:
        return jsonify({'message': 'Unbale to connect to the server. Make sure the server is running', 'error': str(e)}), 500


# Get video_tasks route
@app.route('/v1/video_tasks/<prompt_id>', methods=['GET'])
def video_tasks(prompt_id):
    token = request.headers.get('Authorization')
    if not token or not token.startswith("Bearer "):
        return jsonify({'error': 'Valid Bearer token required'}), 400
    token = token.split(" ")[1]

    try:

        # Establish WebSocket connection to fetch real-time status
        ws = websocket.WebSocket()
        ws.connect(f"{ws_address}?clientId={client_id}&token={token}",
                   header={"Authorization": f"Bearer {token}"})

        # Request current status of the specific prompt
        ws.send(json.dumps({"type": "get_status", "prompt_id": prompt_id}))
        response = json.loads(ws.recv())

        # Extract the necessary fields for the specific prompt
        queue_remaining = response.get('data', {}).get('status', {}).get('exec_info', {}).get('queue_remaining', 0)

        # Now proceed to check if the prompt has completed successfully
        history = get_history(prompt_id, token).get(prompt_id, {})

        if not history:
            return jsonify({
                'message': 'Video is being generated.',
                'status': 'pending',
                'prompts_in_queue': queue_remaining
            }), 202

        video_data = None

        # Extract video or GIF details
        for node_id, node_output in history.get('outputs', {}).items():
            if 'gifs' in node_output:
                video = node_output['gifs'][0]  # Take the first available GIF/video

                try:
                    video_data = get_video_data(video['filename'], video['subfolder'], token)
                    break  # Stop after fetching the first valid video
                except Exception as e:
                    print(f"Failed to retrieve video: {str(e)}")

        if not video_data:
            return jsonify({'error': 'Failed to retrieve video data.'}), 500

        # Save the video locally
        # local_video_path = f"static/generated_image_to_video_{prompt_id}.mp4"
        # with open(local_video_path, 'wb') as f:
        #     f.write(video_data)

        # Send the video as an HTTP response
        return send_file(
            io.BytesIO(video_data),
            mimetype='video/mp4',
            as_attachment=True,
            download_name='generated_video.mp4'
        )

    except Exception as e:
        return jsonify({'error': str(e)}), 500

# Route: Image to Video old
@app.route('/image_to_video', methods=['POST'])
def image_to_video():
    data = request.json

    # Extract and validate token
    token = request.headers.get('Authorization')
    if not token or not token.startswith("Bearer "):
        return jsonify({'error': 'Valid Bearer token required'}), 400
    token = token.split(" ")[1]

    # Validate text prompt
    text_prompt = data.get('text_prompt')
    frame_rate = data.get('frame_rate')
    steps = data.get('steps')
    if not text_prompt or not text_prompt.strip():
        return jsonify({'error': 'Text prompt is required'}), 400

    # Check if frame_rate is missing or invalid
    if not frame_rate:  # If frame_rate is None, empty, or 0
        frame_rate = 24  # Default to 24 fps
    else:
        try:
            frame_rate = int(frame_rate)
            if frame_rate not in [8, 12, 24]:
                return jsonify({'error': 'Frame rate must be a valid number (8, 12, or 24).'}), 400
        except ValueError:
            return jsonify({'error': 'Frame rate must be a valid number (8, 12, or 24).'}), 400

    if not steps:
        steps = 50

    # Handle uploaded image or base64 image
    image_file = request.files.get('image')
    base64_image = data.get('base64_image')

    image_path = None  # Initialize image path

    try:
        if image_file:
            # Check if the file has an allowed extension
            if not allowed_file(image_file.filename):
                return jsonify({'error': 'Unsupported image format'}), 400

            # Secure the filename
            filename = secure_filename(image_file.filename)

            # Generate a unique path for the image
            unique_filename = f"{uuid.uuid4()}_{filename}"
            image_path = os.path.join('static', unique_filename)

            # Ensure the 'static' directory exists
            os.makedirs('static', exist_ok=True)

            # Save the image to the static directory
            image_file.save(image_path)

            # Construct the public URL to access the image
            image_url = f"https://gosign-de-comfyui-api.hf.space/{image_path}"

        elif base64_image:
            # Save base64 image
            try:
                image_path, image_url = save_base64_image(base64_image)
            except Exception as e:
                raise ValueError(f'Invalid base64 image data: {str(e)}')

        else:
            return jsonify({'error': 'Image is required (file or base64)'}), 400

        # Load workflow configuration
        current_dir = os.path.dirname(os.path.abspath(__file__))
        workflow_path = os.path.join(current_dir, 'workflows/cogvideox_image_to_video_workflow_api.json')
        with open(workflow_path, 'r', encoding='utf-8') as f:
            workflow = json.load(f)

        # Modify workflow with inputs
        workflow["30"]["inputs"]["prompt"] = text_prompt
        workflow["31"]["inputs"]["prompt"] = "Low quality, watermark, strange motion, blur"
        workflow["44"]["inputs"]["frame_rate"] = frame_rate
        workflow["57"]["inputs"]["steps"] = steps
        workflow["73"]["inputs"]["url"] = image_url

        # WebSocket connection to queue and monitor workflow
        ws = websocket.WebSocket()
        ws.connect(f"{ws_address}?clientId={client_id}&token={token}",
                   header={"Authorization": f"Bearer {token}"})

        # Queue the prompt and wait for completion
        prompt_id = queue_prompt(workflow, token)


        # Wait for workflow execution to complete
        last_ping = time.time()
        PING_INTERVAL = 30

        # Wait for workflow execution to complete
        while True:
            message = json.loads(ws.recv())
            if message.get('type') == 'executing' and message['data']['node'] is None \
                    and message['data']['prompt_id'] == prompt_id:
                break

            # Send a ping if PING_INTERVAL has passed
            if time.time() - last_ping > PING_INTERVAL:
                ws.send('ping')
                last_ping = time.time()

        # Fetch the history of the workflow
        history = get_history(prompt_id, token).get(prompt_id, {})
        video_data = None

        # Find the video or GIF data from the outputs
        for node_id, node_output in history.get('outputs', {}).items():
            if 'gifs' in node_output:
                video = node_output['gifs'][0]  # Take the first video/GIF
                try:
                    video_data = get_video_data(video['filename'], video['subfolder'], token)
                    break  # Stop after fetching the first valid video
                except Exception as e:
                    print(f"Failed to retrieve video: {str(e)}")

        # Ensure video data was retrieved
        if not video_data:
            raise ValueError('Failed to generate video')

        # Save the video locally
        local_video_path = f"static/generated_image_to_video_{prompt_id}.mp4"
        with open(local_video_path, 'wb') as f:
            f.write(video_data)

        # Construct the public URL for the video

        # video_url = f"https://gosign-de-comfyui-api.hf.space/{local_video_path}"

        # Prepare the response with the video URL
        # response_data = {
        #     'video_url': video_url,
        #     'message': 'Video generated successfully'
        # }

        # return jsonify(response_data), 200

        # Send the video as an HTTP response
        response = send_file(
            io.BytesIO(video_data),
            mimetype='video/mp4',
            as_attachment=True,
            download_name='generated_video.mp4'
        )

        return response

    except Exception as e:
        return jsonify({'error': str(e)}), 500

    finally:
        # Always delete the image if it was saved
        if image_path and os.path.exists(image_path):
            os.remove(image_path)
            print(f"Deleted temporary image: {image_path}", flush=True)


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)