File size: 32,062 Bytes
4f65198
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c8b357c
4f65198
 
 
 
 
 
 
 
 
 
151425f
 
 
 
fb42ffd
89fcc58
 
 
151425f
 
 
 
 
 
25a86d2
 
 
 
 
 
151425f
 
 
 
25a86d2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
151425f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f65198
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f627e1e
4f65198
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
756bd1b
4f65198
8458849
b177918
 
 
 
 
 
 
 
 
 
 
 
 
 
756bd1b
 
b177918
4f65198
b177918
756bd1b
 
 
4f65198
 
 
 
b177918
4f65198
 
 
 
 
 
 
 
 
 
b177918
4f65198
 
 
 
 
8458849
 
4f65198
 
 
 
 
b177918
4f65198
 
 
 
 
8458849
4f65198
 
b177918
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f65198
 
 
 
8458849
4f65198
 
 
 
 
 
 
 
b177918
 
 
4f65198
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0de0de4
 
 
 
 
38ded30
1f82bf6
f101e49
 
 
 
83e5933
 
f101e49
0bd1248
8619bbc
 
 
 
 
 
 
 
0bd1248
 
8619bbc
f101e49
8619bbc
 
 
0bd1248
 
8619bbc
 
 
 
f101e49
8619bbc
f101e49
 
 
8619bbc
f101e49
0bd1248
f101e49
 
 
8619bbc
f101e49
 
8619bbc
f101e49
 
8619bbc
 
 
f101e49
8619bbc
f101e49
0bd1248
8619bbc
f101e49
8619bbc
0bd1248
f101e49
 
 
 
 
8619bbc
f101e49
8619bbc
f101e49
0bd1248
8619bbc
 
 
f101e49
0bd1248
 
8619bbc
72ef525
83e5933
f101e49
 
8619bbc
 
 
 
 
 
4f65198
83e5933
 
c0042d0
 
83e5933
c0042d0
 
 
 
 
 
83e5933
c0042d0
 
 
83e5933
 
52e7e4c
 
 
 
 
 
 
 
 
83e5933
52e7e4c
 
bd95f41
d940931
83e5933
d940931
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f65198
83e5933
 
c7ce94a
9c6e21f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3764c21
9c6e21f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44061da
9c6e21f
44061da
9c6e21f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83e5933
 
 
4f65198
83e5933
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
import os
import re
import random
from http import HTTPStatus
from typing import Dict, List, Optional, Tuple
import base64
import anthropic
import openai
import asyncio
import time
from functools import partial
import json
import gradio as gr
import modelscope_studio.components.base as ms
import modelscope_studio.components.legacy as legacy
import modelscope_studio.components.antd as antd
import html
import urllib.parse
from huggingface_hub import HfApi, create_repo, hf_hub_download
import string
import requests
from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.common.exceptions import WebDriverException, TimeoutException
from PIL import Image
from io import BytesIO
from datetime import datetime
import spaces
from safetensors.torch import load_file
from diffusers import FluxPipeline
import torch
from os import path  # ์ด ์ค„์„ ์ถ”๊ฐ€
from datetime import datetime, timedelta
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
# ์บ์‹œ ๊ฒฝ๋กœ ์„ค์ •
cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
os.environ["TRANSFORMERS_CACHE"] = cache_path
os.environ["HF_HUB_CACHE"] = cache_path
os.environ["HF_HOME"] = cache_path


# Hugging Face ํ† ํฐ ์„ค์ •
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    print("Warning: HF_TOKEN not found in environment variables")

# FLUX ๋ชจ๋ธ ์ดˆ๊ธฐํ™”
if not path.exists(cache_path):
    os.makedirs(cache_path, exist_ok=True)

try:
    pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev",
        torch_dtype=torch.bfloat16,
        use_auth_token=HF_TOKEN  # Hugging Face ํ† ํฐ ์ถ”๊ฐ€
    )
    pipe.load_lora_weights(
        hf_hub_download(
            "ByteDance/Hyper-SD",
            "Hyper-FLUX.1-dev-8steps-lora.safetensors",
            token=HF_TOKEN  # Hugging Face ํ† ํฐ ์ถ”๊ฐ€
        )
    )
    pipe.fuse_lora(lora_scale=0.125)
    pipe.to(device="cuda", dtype=torch.bfloat16)
    print("Successfully initialized FLUX model with authentication")
except Exception as e:
    print(f"Error initializing FLUX model: {str(e)}")
    pipe = None
    


# ์ด๋ฏธ์ง€ ์ƒ์„ฑ ํ•จ์ˆ˜ ์ถ”๊ฐ€
@spaces.GPU
def generate_image(prompt, height=512, width=512, steps=8, scales=3.5, seed=3413):
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
        return pipe(
            prompt=[prompt],
            generator=torch.Generator().manual_seed(int(seed)),
            num_inference_steps=int(steps),
            guidance_scale=float(scales),
            height=int(height),
            width=int(width),
            max_sequence_length=256
        ).images[0]
        
# SystemPrompt ๋ถ€๋ถ„์„ ์ง์ ‘ ์ •์˜
SystemPrompt = """You are 'MOUSE-I', an advanced AI visualization expert. Your mission is to transform every response into a visually stunning and highly informative presentation.

Core Capabilities:
- Transform text responses into rich visual experiences
- Create interactive data visualizations and charts
- Design beautiful and intuitive user interfaces
- Utilize engaging animations and transitions
- Present information in a clear, structured manner

Visual Elements to Include:
- Charts & Graphs (using Chart.js, D3.js)
- Interactive Data Visualizations
- Modern UI Components
- Engaging Animations
- Informative Icons & Emojis
- Color-coded Information Blocks
- Progress Indicators
- Timeline Visualizations
- Statistical Representations
- Comparison Tables

Technical Requirements:
- Modern HTML5/CSS3/JavaScript
- Responsive Design
- Interactive Elements
- Clean Typography
- Professional Color Schemes
- Smooth Animations
- Cross-browser Compatibility

Libraries Available:
- Chart.js for Data Visualization
- D3.js for Complex Graphics
- Bootstrap for Layout
- jQuery for Interactions
- Three.js for 3D Elements

Design Principles:
- Visual Hierarchy
- Clear Information Flow
- Consistent Styling
- Intuitive Navigation
- Engaging User Experience
- Accessibility Compliance

Remember to:
- Present data in the most visually appealing way
- Use appropriate charts for different data types
- Include interactive elements where relevant
- Maintain a professional and modern aesthetic
- Ensure responsive design for all devices

Return only HTML code wrapped in code blocks, focusing on creating visually stunning and informative presentations.
"""

from config import DEMO_LIST

class Role:
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"

History = List[Tuple[str, str]]
Messages = List[Dict[str, str]]

# ์ด๋ฏธ์ง€ ์บ์‹œ๋ฅผ ๋ฉ”๋ชจ๋ฆฌ์— ์ €์žฅ
IMAGE_CACHE = {}

# boost_prompt ํ•จ์ˆ˜์™€ handle_boost ํ•จ์ˆ˜๋ฅผ ์ถ”๊ฐ€ํ•ฉ๋‹ˆ๋‹ค
def boost_prompt(prompt: str) -> str:
    if not prompt:
        return ""
    
    # ์ฆ๊ฐ•์„ ์œ„ํ•œ ์‹œ์Šคํ…œ ํ”„๋กฌํ”„ํŠธ
    boost_system_prompt = """
    ๋‹น์‹ ์€ ์›น ๊ฐœ๋ฐœ ํ”„๋กฌํ”„ํŠธ ์ „๋ฌธ๊ฐ€์ž…๋‹ˆ๋‹ค. 
    ์ฃผ์–ด์ง„ ํ”„๋กฌํ”„ํŠธ๋ฅผ ๋ถ„์„ํ•˜์—ฌ ๋” ์ƒ์„ธํ•˜๊ณ  ์ „๋ฌธ์ ์ธ ์š”๊ตฌ์‚ฌํ•ญ์œผ๋กœ ํ™•์žฅํ•˜๋˜,
    ์›๋ž˜ ์˜๋„์™€ ๋ชฉ์ ์€ ๊ทธ๋Œ€๋กœ ์œ ์ง€ํ•˜๋ฉด์„œ ๋‹ค์Œ ๊ด€์ ๋“ค์„ ๊ณ ๋ คํ•˜์—ฌ ์ฆ๊ฐ•ํ•˜์‹ญ์‹œ์˜ค:
    1. ๊ธฐ์ˆ ์  ๊ตฌํ˜„ ์ƒ์„ธ
    2. UI/UX ๋””์ž์ธ ์š”์†Œ
    3. ์‚ฌ์šฉ์ž ๊ฒฝํ—˜ ์ตœ์ ํ™”
    4. ์„ฑ๋Šฅ๊ณผ ๋ณด์•ˆ
    5. ์ ‘๊ทผ์„ฑ๊ณผ ํ˜ธํ™˜์„ฑ
    
    ๊ธฐ์กด SystemPrompt์˜ ๋ชจ๋“  ๊ทœ์น™์„ ์ค€์ˆ˜ํ•˜๋ฉด์„œ ์ฆ๊ฐ•๋œ ํ”„๋กฌํ”„ํŠธ๋ฅผ ์ƒ์„ฑํ•˜์‹ญ์‹œ์˜ค.
    """
    
    try:
        # Claude API ์‹œ๋„
        try:
            response = claude_client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=2000,
                messages=[{
                    "role": "user",
                    "content": f"๋‹ค์Œ ํ”„๋กฌํ”„ํŠธ๋ฅผ ๋ถ„์„ํ•˜๊ณ  ์ฆ๊ฐ•ํ•˜์‹œ์˜ค: {prompt}"
                }]
            )
            
            if hasattr(response, 'content') and len(response.content) > 0:
                return response.content[0].text
            raise Exception("Claude API ์‘๋‹ต ํ˜•์‹ ์˜ค๋ฅ˜")
            
        except Exception as claude_error:
            print(f"Claude API ์—๋Ÿฌ, OpenAI๋กœ ์ „ํ™˜: {str(claude_error)}")
            
            # OpenAI API ์‹œ๋„
            completion = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": boost_system_prompt},
                    {"role": "user", "content": f"๋‹ค์Œ ํ”„๋กฌํ”„ํŠธ๋ฅผ ๋ถ„์„ํ•˜๊ณ  ์ฆ๊ฐ•ํ•˜์‹œ์˜ค: {prompt}"}
                ],
                max_tokens=2000,
                temperature=0.7
            )
            
            if completion.choices and len(completion.choices) > 0:
                return completion.choices[0].message.content
            raise Exception("OpenAI API ์‘๋‹ต ํ˜•์‹ ์˜ค๋ฅ˜")
            
    except Exception as e:
        print(f"ํ”„๋กฌํ”„ํŠธ ์ฆ๊ฐ• ์ค‘ ์˜ค๋ฅ˜ ๋ฐœ์ƒ: {str(e)}")
        return prompt  # ์˜ค๋ฅ˜ ๋ฐœ์ƒ์‹œ ์›๋ณธ ํ”„๋กฌํ”„ํŠธ ๋ฐ˜ํ™˜

# Boost ๋ฒ„ํŠผ ์ด๋ฒคํŠธ ํ•ธ๋“ค๋Ÿฌ
def handle_boost(prompt: str):
    try:
        boosted_prompt = boost_prompt(prompt)
        return boosted_prompt, gr.update(active_key="empty")
    except Exception as e:
        print(f"Boost ์ฒ˜๋ฆฌ ์ค‘ ์˜ค๋ฅ˜: {str(e)}")
        return prompt, gr.update(active_key="empty")
        
def get_image_base64(image_path):
    if image_path in IMAGE_CACHE:
        return IMAGE_CACHE[image_path]
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode()
            IMAGE_CACHE[image_path] = encoded_string
            return encoded_string
    except:
        return IMAGE_CACHE.get('default.png', '')
        
def history_to_messages(history: History, system: str) -> Messages:
    messages = [{'role': Role.SYSTEM, 'content': system}]
    for h in history:
        messages.append({'role': Role.USER, 'content': h[0]})
        messages.append({'role': Role.ASSISTANT, 'content': h[1]})
    return messages

def messages_to_history(messages: Messages) -> History:
    assert messages[0]['role'] == Role.SYSTEM
    history = []
    for q, r in zip(messages[1::2], messages[2::2]):
        history.append([q['content'], r['content']])
    return history

# API ํด๋ผ์ด์–ธํŠธ ์ดˆ๊ธฐํ™”
YOUR_ANTHROPIC_TOKEN = os.getenv('ANTHROPIC_API_KEY', '')  # ๊ธฐ๋ณธ๊ฐ’ ์ถ”๊ฐ€
YOUR_OPENAI_TOKEN = os.getenv('OPENAI_API_KEY', '')  # ๊ธฐ๋ณธ๊ฐ’ ์ถ”๊ฐ€

# API ํ‚ค ๊ฒ€์ฆ
if not YOUR_ANTHROPIC_TOKEN or not YOUR_OPENAI_TOKEN:
    print("Warning: API keys not found in environment variables")

# API ํด๋ผ์ด์–ธํŠธ ์ดˆ๊ธฐํ™” ์‹œ ์˜ˆ์™ธ ์ฒ˜๋ฆฌ ์ถ”๊ฐ€
try:
    claude_client = anthropic.Anthropic(api_key=YOUR_ANTHROPIC_TOKEN)
    openai_client = openai.OpenAI(api_key=YOUR_OPENAI_TOKEN)
except Exception as e:
    print(f"Error initializing API clients: {str(e)}")
    claude_client = None
    openai_client = None

# try_claude_api ํ•จ์ˆ˜ ์ˆ˜์ •
async def try_claude_api(system_message, claude_messages, timeout=15):
    try:
        start_time = time.time()
        with claude_client.messages.stream(
            model="claude-3-5-sonnet-20241022",
            max_tokens=7860,
            system=system_message,
            messages=claude_messages
        ) as stream:
            collected_content = ""
            for chunk in stream:
                current_time = time.time()
                if current_time - start_time > timeout:
                    print(f"Claude API response time: {current_time - start_time:.2f} seconds")
                    raise TimeoutError("Claude API timeout")
                if chunk.type == "content_block_delta":
                    collected_content += chunk.delta.text
                    yield collected_content
                    await asyncio.sleep(0)
                
                start_time = current_time
                
    except Exception as e:
        print(f"Claude API error: {str(e)}")
        raise e

async def try_openai_api(openai_messages):
    try:
        stream = openai_client.chat.completions.create(
            model="gpt-4o",
            messages=openai_messages,
            stream=True,
            max_tokens=4096,
            temperature=0.7
        )
        
        collected_content = ""
        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                collected_content += chunk.choices[0].delta.content
                yield collected_content
                
    except Exception as e:
        print(f"OpenAI API error: {str(e)}")
        raise e

class Demo:
    def __init__(self):
        pass

    async def generation_code(self, query: Optional[str], _setting: Dict[str, str]):
        if not query or query.strip() == '':
            query = get_random_placeholder()
        
        # ์ด๋ฏธ์ง€ ์ƒ์„ฑ์ด ํ•„์š”ํ•œ์ง€ ํ™•์ธ
        needs_image = '์ด๋ฏธ์ง€' in query or '๊ทธ๋ฆผ' in query or 'image' in query.lower()
        image_prompt = None
        
        # ์ด๋ฏธ์ง€ ํ”„๋กฌํ”„ํŠธ ์ถ”์ถœ
        if needs_image:
            for keyword in ['์ด๋ฏธ์ง€:', '๊ทธ๋ฆผ:', 'image:']:
                if keyword in query.lower():
                    image_prompt = query.split(keyword)[1].strip()
                    break
            if not image_prompt:
                image_prompt = query  # ๋ช…์‹œ์  ํ”„๋กฌํ”„ํŠธ๊ฐ€ ์—†์œผ๋ฉด ์ „์ฒด ์ฟผ๋ฆฌ ์‚ฌ์šฉ
        
        messages = [{'role': Role.SYSTEM, 'content': _setting['system']}]
        messages.append({'role': Role.USER, 'content': query})
        
        system_message = messages[0]['content']
        claude_messages = [{"role": "user", "content": query}]
        openai_messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": query}
        ]
        
        try:
            yield [
                "",
                None,
                gr.update(active_key="loading"),
                gr.update(open=True)
            ]
            await asyncio.sleep(0)

            collected_content = None
            try:
                async for content in try_claude_api(system_message, claude_messages):
                    yield [
                        "",
                        None,
                        gr.update(active_key="loading"),
                        gr.update(open=True)
                    ]
                    await asyncio.sleep(0)
                    collected_content = content
                    
            except Exception as claude_error:
                print(f"Falling back to OpenAI API due to Claude error: {str(claude_error)}")
                
                async for content in try_openai_api(openai_messages):
                    yield [
                        "",
                        None,
                        gr.update(active_key="loading"),
                        gr.update(open=True)
                    ]
                    await asyncio.sleep(0)
                    collected_content = content

            if collected_content:
                # ์ด๋ฏธ์ง€ ์ƒ์„ฑ์ด ํ•„์š”ํ•œ ๊ฒฝ์šฐ
                if needs_image and image_prompt:
                    try:
                        print(f"Generating image for prompt: {image_prompt}")
                        # FLUX ๋ชจ๋ธ์„ ์‚ฌ์šฉํ•˜์—ฌ ์ด๋ฏธ์ง€ ์ƒ์„ฑ
                        if pipe is not None:
                            image = generate_image(
                                prompt=image_prompt,
                                height=512,
                                width=512,
                                steps=8,
                                scales=3.5,
                                seed=random.randint(1, 10000)
                            )
                            
                            # ์ด๋ฏธ์ง€๋ฅผ Base64๋กœ ์ธ์ฝ”๋”ฉ
                            buffered = BytesIO()
                            image.save(buffered, format="PNG")
                            img_str = base64.b64encode(buffered.getvalue()).decode()
                            
                            # HTML์— ์ด๋ฏธ์ง€ ์ถ”๊ฐ€
                            image_html = f'''
                            <div class="generated-image" style="margin: 20px 0; text-align: center;">
                                <h3 style="color: #333; margin-bottom: 10px;">Generated Image:</h3>
                                <img src="data:image/png;base64,{img_str}" 
                                     style="max-width: 100%; 
                                            border-radius: 10px; 
                                            box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
                                <p style="color: #666; margin-top: 10px; font-style: italic;">
                                    Prompt: {html.escape(image_prompt)}
                                </p>
                            </div>
                            '''
                            
                            # HTML ์‘๋‹ต์— ์ด๋ฏธ์ง€ ์‚ฝ์ž…
                            if '```html' in collected_content:
                                # HTML ์ฝ”๋“œ ๋ธ”๋ก ๋‚ด๋ถ€์— ์ด๋ฏธ์ง€ ์ถ”๊ฐ€
                                collected_content = collected_content.replace('```html\n', f'```html\n{image_html}')
                            else:
                                # HTML ์ฝ”๋“œ ๋ธ”๋ก์œผ๋กœ ๊ฐ์‹ธ์„œ ์ด๋ฏธ์ง€ ์ถ”๊ฐ€
                                collected_content = f'```html\n{image_html}\n```\n{collected_content}'
                            
                            print("Image generation successful")
                        else:
                            raise Exception("FLUX model not initialized")
                            
                    except Exception as e:
                        print(f"Image generation error: {str(e)}")
                        error_message = f'''
                        <div style="color: #ff4d4f; padding: 10px; margin: 10px 0; 
                                    border-left: 4px solid #ff4d4f; background: #fff2f0;">
                            <p>Failed to generate image: {str(e)}</p>
                        </div>
                        '''
                        if '```html' in collected_content:
                            collected_content = collected_content.replace('```html\n', f'```html\n{error_message}')
                        else:
                            collected_content = f'```html\n{error_message}\n```\n{collected_content}'

                # ์ตœ์ข… ๊ฒฐ๊ณผ ํ‘œ์‹œ
                yield [
                    collected_content,
                    send_to_sandbox(remove_code_block(collected_content)),
                    gr.update(active_key="render"),
                    gr.update(open=False)
                ]
            else:
                raise ValueError("No content was generated from either API")
                
        except Exception as e:
            print(f"Error details: {str(e)}")
            raise ValueError(f'Error calling APIs: {str(e)}')

    def clear_history(self):
        return []

def remove_code_block(text):
    pattern = r'```html\n(.+?)\n```'
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    else:
        return text.strip()

def history_render(history: History):
    return gr.update(open=True), history

def send_to_sandbox(code):
    encoded_html = base64.b64encode(code.encode('utf-8')).decode('utf-8')
    data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
    return f"""
        <iframe 
            src="{data_uri}" 
            style="width:100%; height:800px; border:none;"
            frameborder="0"
        ></iframe>
    """
# ๋ฐฐํฌ ๊ด€๋ จ ํ•จ์ˆ˜ ์ถ”๊ฐ€
def generate_space_name():
    """6์ž๋ฆฌ ๋žœ๋ค ์˜๋ฌธ ์ด๋ฆ„ ์ƒ์„ฑ"""
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for i in range(6))

def deploy_to_vercel(code: str):
    try:
        token = "A8IFZmgW2cqA4yUNlLPnci0N"
        if not token:
            return "Vercel ํ† ํฐ์ด ์„ค์ •๋˜์ง€ ์•Š์•˜์Šต๋‹ˆ๋‹ค."

        # 6์ž๋ฆฌ ์˜๋ฌธ ํ”„๋กœ์ ํŠธ ์ด๋ฆ„ ์ƒ์„ฑ
        project_name = ''.join(random.choice(string.ascii_lowercase) for i in range(6))


        # Vercel API ์—”๋“œํฌ์ธํŠธ
        deploy_url = "https://api.vercel.com/v13/deployments"

        # ํ—ค๋” ์„ค์ •
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        # package.json ํŒŒ์ผ ์ƒ์„ฑ
        package_json = {
            "name": project_name,
            "version": "1.0.0",
            "private": True,  # true -> True๋กœ ์ˆ˜์ •
            "dependencies": {
                "vite": "^5.0.0"
            },
            "scripts": {
                "dev": "vite",
                "build": "echo 'No build needed' && mkdir -p dist && cp index.html dist/",
                "preview": "vite preview"
            }
        }

        # ๋ฐฐํฌํ•  ํŒŒ์ผ ๋ฐ์ดํ„ฐ ๊ตฌ์กฐ
        files = [
            {
                "file": "index.html",
                "data": code
            },
            {
                "file": "package.json",
                "data": json.dumps(package_json, indent=2)  # indent ์ถ”๊ฐ€๋กœ ๊ฐ€๋…์„ฑ ํ–ฅ์ƒ
            }
        ]

        # ํ”„๋กœ์ ํŠธ ์„ค์ •
        project_settings = {
            "buildCommand": "npm run build",
            "outputDirectory": "dist",
            "installCommand": "npm install",
            "framework": None
        }

        # ๋ฐฐํฌ ์š”์ฒญ ๋ฐ์ดํ„ฐ
        deploy_data = {
            "name": project_name,
            "files": files,
            "target": "production",
            "projectSettings": project_settings
        }


        deploy_response = requests.post(deploy_url, headers=headers, json=deploy_data)
        
        if deploy_response.status_code != 200:
            return f"๋ฐฐํฌ ์‹คํŒจ: {deploy_response.text}"

        # URL ํ˜•์‹ ์ˆ˜์ • - 6์ž๋ฆฌ.vercel.app ํ˜•ํƒœ๋กœ ๋ฐ˜ํ™˜
        deployment_url = f"{project_name}.vercel.app"
        
        time.sleep(5)

        return f"""๋ฐฐํฌ ์™„๋ฃŒ! <a href="https://{deployment_url}" target="_blank" style="color: #1890ff; text-decoration: underline; cursor: pointer;">https://{deployment_url}</a>"""

    except Exception as e:
        return f"๋ฐฐํฌ ์ค‘ ์˜ค๋ฅ˜ ๋ฐœ์ƒ: {str(e)}"
        
theme = gr.themes.Soft()

def get_random_placeholder():
    return random.choice(DEMO_LIST)['description']

def update_placeholder():
    return gr.update(placeholder=get_random_placeholder())


def create_main_interface():
     """๋ฉ”์ธ ์ธํ„ฐํŽ˜์ด์Šค ์ƒ์„ฑ ํ•จ์ˆ˜"""
     
     #NEW - ๊ฒ€์ƒ‰ ๊ฒฐ๊ณผ๋ฅผ ํ†ตํ•ฉํ•œ ์‘๋‹ต ์ƒ์„ฑ ํ•จ์ˆ˜
     async def execute_search_and_generate(query, setting):
          try:
               print(f"Executing search for query: {query}")
               
               # ๊ฒ€์ƒ‰ ์‹คํ–‰
               url = "https://api.serphouse.com/serp/live"
               payload = {
                    "data": {
                         "q": query,
                         "domain": "google.com", 
                         "lang": "en",
                         "device": "desktop",
                         "serp_type": "news",
                         "loc": "United States",
                         "page": "1",
                         "num": "10"
                    }
               }
               headers = {
                    "Authorization": "Bearer V38CNn4HXpLtynJQyOeoUensTEYoFy8PBUxKpDqAW1pawT1vfJ2BWtPQ98h6",
                    "Content-Type": "application/json"
               }
               
               response = requests.post(url, headers=headers, json=payload)
               results = response.json()
               print(f"Search results: {results}")  # ๋””๋ฒ„๊น…์šฉ
               
               # ๊ฒ€์ƒ‰ ๊ฒฐ๊ณผ๋ฅผ HTML๋กœ ๋ณ€ํ™˜
               search_content = "```html\n<div class='search-results'>\n"
               search_content += "<h2>์ตœ์‹  ๋‰ด์Šค ๊ฒ€์ƒ‰ ๊ฒฐ๊ณผ</h2>\n"
               
               # API ์‘๋‹ต ๊ตฌ์กฐ์— ๋งž๊ฒŒ ์ˆ˜์ •
               if 'results' in results:
                    news_items = results['results'].get('news', [])
                    for item in news_items[:5]:
                         search_content += f"""
                              <div class="search-item">
                                   <h3><a href="{item['url']}" target="_blank">{item['title']}</a></h3>
                                   <p>{item['snippet']}</p>
                                   <div class="search-meta">
                                        <span class="source">{item['channel']}</span>
                                        <span class="time">{item['time']}</span>
                                   </div>
                              </div>
                         """
               search_content += "</div>\n```"

               # ๊ฒ€์ƒ‰ ๊ฒฐ๊ณผ๋ฅผ ํฌํ•จํ•œ ํ”„๋กฌํ”„ํŠธ ์ƒ์„ฑ
               enhanced_prompt = f"""Based on these news search results, create a comprehensive visual summary:

{search_content}

Please create a visually appealing HTML response that:
1. Summarizes the key points from the news
2. Organizes information in a clear structure
3. Uses appropriate HTML formatting and styling
4. Includes relevant quotes and statistics
5. Provides proper source attribution

The response should be in HTML format with appropriate styling."""

               print("Generating response with search results...")  # ๋””๋ฒ„๊น…์šฉ
               
               # async generator๋ฅผ ์ฒ˜๋ฆฌํ•˜๊ธฐ ์œ„ํ•œ ์ˆ˜์ •
               async for result in demo_instance.generation_code(enhanced_prompt, setting):
                    final_result = result
                    print(f"Generated result: {final_result}")  # ๋””๋ฒ„๊น…์šฉ
               
               print("Response generation completed")  # ๋””๋ฒ„๊น…์šฉ
               return final_result
                    
          except Exception as e:
               print(f"Search error: {str(e)}")
               print(f"Full error details: {str(e.__class__.__name__)}: {str(e)}")
               return [
                    "",
                    None,
                    gr.update(active_key="error"),
                    gr.update(open=False)
               ]

     def execute_code(query: str):
          if not query or query.strip() == '':
               return None, gr.update(active_key="empty")
          
          try:
               if '```html' in query and '```' in query:
                    code = remove_code_block(query)
               else:
                    code = query.strip()
               
               return send_to_sandbox(code), gr.update(active_key="render")
          except Exception as e:
               print(f"Error executing code: {str(e)}")
               return None, gr.update(active_key="empty")

     async def handle_generation(query, setting, is_search):
          try:
               print(f"Mode: {'Web Search' if is_search else 'Generate'}")  # ๋””๋ฒ„๊น…์šฉ
               if is_search:
                    print("Executing search and generate...")  # ๋””๋ฒ„๊น…์šฉ
                    return await execute_search_and_generate(query, setting)
               else:
                    print("Executing normal generation...")  # ๋””๋ฒ„๊น…์šฉ
                    async for result in demo_instance.generation_code(query, setting):
                         final_result = result
                    return final_result
          except Exception as e:
               print(f"Generation error: {str(e)}")
               return ["", None, gr.update(active_key="error"), gr.update(open=False)]

     # CSS ํŒŒ์ผ ๋‚ด์šฉ์„ ์ง์ ‘ ์ ์šฉ
     with open('app.css', 'r', encoding='utf-8') as f:
          custom_css = f.read()

     #NEW - ๊ฒ€์ƒ‰ ๊ฒฐ๊ณผ ์Šคํƒ€์ผ ์ถ”๊ฐ€
     custom_css += """
          .search-summary {
               margin: 20px 0;
               padding: 20px;
               background: #f8f9fa;
               border-radius: 10px;
          }

          .search-item {
               margin-bottom: 15px;
               padding: 15px;
               border-left: 4px solid #007aff;
               background: white;
               border-radius: 4px;
          }

          .search-item h3 {
               margin: 0 0 10px 0;
               color: #1a0dab;
          }

          .search-item a {
               color: inherit;
               text-decoration: none;
          }

          .search-item p {
               margin: 0;
               color: #4d5156;
               font-size: 14px;
               line-height: 1.5;
          }

          .empty-content {
               padding: 40px !important;
               background: #f8f9fa !important;
               border-radius: 10px !important;
               margin: 20px !important;
          }    
     
          .container {
               background: #f0f0f0;
               min-height: 100vh;
               padding: 20px;
               display: flex;
               justify-content: center;
               align-items: center;
               font-family: -apple-system, BlinkMacSystemFont, sans-serif;
          }

          .mode-selector {
               margin-bottom: 15px;
               padding: 10px;
               border-radius: 8px;
               background: #f8f9fa;
          }

          .error-content {
               padding: 20px;
               background: #fff2f0;
               border-radius: 8px;
               border: 1px solid #ff4d4f;
          }

          .search-mode-active {
               background-color: rgba(66,133,244,0.1);
          }
     """

     demo = gr.Blocks(css=custom_css, theme=theme)

     with demo:
          with gr.Row():  # ์ „์ฒด๋ฅผ ๊ฐ์‹ธ๋Š” Row ์ถ”๊ฐ€
               # ์ขŒ์ธก ํŒจ๋„
               with gr.Column(scale=1):
                    mode = gr.Radio(
                         choices=["Generate", "Generate + Web Search"],
                         label="Mode",
                         value="Generate",
                         info="Select 'Generate + Web Search' to include web search results",
                         elem_classes="mode-selector"
                    )

                    input = gr.Textbox(
                         label="Input",
                         placeholder=get_random_placeholder(),
                         lines=5,
                         elem_classes="custom-textarea"
                    )

                    with gr.Row():
                         btn = gr.Button("Generate", elem_classes="generate-btn")
                         boost_btn = gr.Button("Enhance", elem_classes="enhance-btn")
                         deploy_btn = gr.Button("Share", elem_classes="share-btn")

                    deploy_result = gr.HTML(label="Share Result", elem_classes="deploy-result")

               # ์šฐ์ธก ํŒจ๋„
               with gr.Column(scale=2):
                    with gr.Box():  # Box๋กœ ๊ฐ์‹ธ์„œ ์‹œ๊ฐ์  ๊ตฌ๋ถ„
                         gr.HTML("""
                              <div class="window-frame">
                                   <div class="window-header">
                                        <div class="window-controls">
                                             <div class="control close"></div>
                                             <div class="control minimize"></div>
                                             <div class="control maximize"></div>
                                        </div>
                                        <div class="window-title">Preview</div>
                                   </div>
                              </div>
                         """)
                         
                         # ๊ฒฐ๊ณผ ํ‘œ์‹œ ์˜์—ญ
                         with gr.Tabs(selected="empty") as state_tab:
                              with gr.TabItem("empty"):
                                   gr.Markdown("Enter your question to begin")
                              with gr.TabItem("loading"):
                                   gr.Markdown("Creating visual presentation...")
                              with gr.TabItem("render"):
                                   sandbox = gr.HTML(elem_classes="html_content")
                              with gr.TabItem("error"):
                                   gr.Markdown("An error occurred. Please try again.")

          # ์ƒํƒœ ๋ณ€์ˆ˜๋“ค
          setting = gr.State({"system": SystemPrompt})
          search_mode = gr.State(False)
          code_output = gr.State("")

          # Drawer ์ปดํฌ๋„ŒํŠธ
          with gr.Box() as code_drawer:
               gr.HTML("""
                    <div class="thinking-container">
                         <!-- thinking container content -->
                    </div>
               """)

          # ์ด๋ฒคํŠธ ํ•ธ๋“ค๋Ÿฌ ์—ฐ๊ฒฐ
          mode.change(fn=lambda x: x == "Generate + Web Search",
                    inputs=[mode],
                    outputs=[search_mode])

          btn.click(fn=handle_generation,
                    inputs=[input, setting, search_mode],
                    outputs=[code_output, sandbox, state_tab, code_drawer])

          boost_btn.click(fn=handle_boost,
                         inputs=[input],
                         outputs=[input, state_tab])

          deploy_btn.click(fn=lambda code: deploy_to_vercel(remove_code_block(code)) if code else "No code to share.",
                         inputs=[code_output],
                         outputs=[deploy_result])

     return demo

if __name__ == "__main__":
     try:
          demo_instance = Demo()
          demo = create_main_interface()
          demo.queue(
               default_concurrency_limit=20,
               status_update_rate=10,
               api_open=False
          ).launch(
               server_name="0.0.0.0",
               server_port=7860,
               share=False,
               debug=False
          )
     except Exception as e:
          print(f"Initialization error: {e}")
          raise