File size: 35,409 Bytes
015dc83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9bd5960
015dc83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d38f1ad
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
015dc83
3c666da
015dc83
 
3c666da
 
 
015dc83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
import spaces
import gradio as gr
import torch                                                                                                                                          
from transformers import AutoTokenizer, AutoModelForCausalLM                                                                                          
import time                                                                                                                                           
import re                                                                                                                                             
import gc
import os
import shutil
from pathlib import Path

# Storage optimization - redirect cache to temporary directories
os.environ["HF_HOME"] = "/tmp/hf_cache"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
os.environ["HF_DATASETS_CACHE"] = "/tmp/datasets_cache"
os.environ["TORCH_HOME"] = "/tmp/torch_cache"

# Venice Edition - NSFW-friendly Dolphin model that fits ZeroGPU comfortably
MODEL_PATH = "dphn/dolphin-2.6-mistral-7b-dpo"

TITLE = """

<div style="text-align: center; margin: 20px 0;">

<h1>🌊 Venice Edition NSFW Enhancer</h1>

<p><strong>Transform JoyCaption output into SEO-optimized, uncensored descriptions</strong></p>

<p><em>Powered by Dolphin 2.6 Mistral 7B DPO - Venice Edition Uncensored</em></p>

</div>

<hr>

"""

print("πŸš€ Loading Venice Edition NSFW Enhancer... v2.1")
print(f"πŸ“¦ Model: {MODEL_PATH}")
print("βœ… Venice Edition uncensored model selected for adult content enhancement")

def cleanup_memory_aggressive():
    """Aggressive memory and cache cleanup with storage optimization"""
    try:
        # Clean up temporary caches to prevent storage overflow
        temp_dirs = ["/tmp/hf_cache", "/tmp/transformers_cache", "/tmp/datasets_cache", "/tmp/torch_cache"]
        for temp_dir in temp_dirs:
            if os.path.exists(temp_dir):
                try:
                    shutil.rmtree(temp_dir, ignore_errors=True)
                except:
                    pass
        
        # Force garbage collection
        gc.collect()
        
        # Clear GPU cache if available
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        
    except Exception as e:
        print(f"⚠️ Cleanup warning: {e}")
        pass

def cleanup_storage():
    """Extended storage cleanup function"""
    cleanup_memory_aggressive()
    print("βœ… Storage and memory cleanup completed")

@spaces.GPU(duration=20)  # Quick title generation
@torch.no_grad()
def generate_title(description_text, keywords=""):
    """Generate catchy, SEO-optimized titles for adult content"""
    
    if not description_text or not description_text.strip():
        return "Please provide description text to generate a title."
    
    start_time = time.time()
    
    try:
        # Create title generation prompt
        system_prompt = "You are a professional adult content title writer. Create catchy, SEO-optimized, clickable titles that are explicit when appropriate and designed to attract viewers while accurately representing the content."
        
        keyword_hint = f" Keywords to consider: {keywords}" if keywords.strip() else ""
        
        user_prompt = f"""Based on this description, create 3 catchy, SEO-optimized titles for adult content:



DESCRIPTION:

{description_text[:500]}

{keyword_hint}



Create titles that are:

- Catchy and clickable

- SEO-optimized for adult content

- Explicit when appropriate

- Under 60 characters each

- Designed to attract viewers



Generate 3 title options:

1.

2. 

3."""
        
        # Format for Dolphin model
        full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        # Tokenize and generate
        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=1200, padding=True)
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=150,
                temperature=0.8,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        # Decode and clean
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        if "<|im_start|>assistant\n" in generated_text:
            titles_text = generated_text.split("<|im_start|>assistant\n")[-1]
        else:
            titles_text = generated_text[len(full_prompt):]
        
        titles_text = titles_text.strip()
        titles_text = re.sub(r'<\|im_.*?\|>', '', titles_text).strip()
        
        # Aggressive cleanup to prevent storage overflow
        del inputs, outputs
        cleanup_memory_aggressive()
        cleanup_storage()
        
        total_time = time.time() - start_time
        return f"βœ… Generated in {total_time:.1f}s\n\n{titles_text}"
        
    except Exception as e:
        cleanup_memory_aggressive()
        return f"Error generating titles: {str(e)[:100]}..."

@spaces.GPU(duration=20)  # Quick keywords generation
@torch.no_grad()
def generate_keywords(seo_keywords_input, description_text="", title=""):
    """Generate input keywords + 3 synonyms for each keyword"""
    
    if not seo_keywords_input or not seo_keywords_input.strip():
        return "Please provide input keywords to generate synonyms."
    
    start_time = time.time()
    
    try:
        # Parse input keywords
        input_keywords = [k.strip() for k in seo_keywords_input.split(',') if k.strip()]
        if not input_keywords:
            return "Please provide valid keywords separated by commas."
        
        # Create keywords generation prompt - EXACTLY as user specified
        system_prompt = "You are a professional SEO specialist for adult content. For each provided keyword, generate exactly 3 relevant synonyms."
        
        context_hint = ""
        if description_text.strip():
            context_hint += f" Context from description: {description_text[:200]}"
        if title.strip():
            context_hint += f" Title: {title[:100]}"
        
        keywords_list = '\n'.join([f"{i+1}. {keyword}" for i, keyword in enumerate(input_keywords)])
        
        user_prompt = f"""For each of these keywords, provide exactly 3 synonyms:



{keywords_list}



Format your response as:

Keyword: synonym1, synonym2, synonym3



{context_hint}



Provide synonyms that are relevant for adult content SEO."""
        
        # Format for Dolphin model
        full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        # Format for Dolphin model
        full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        # Tokenize and generate
        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=1200, padding=True)
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=200,
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        # Decode and clean
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        if "<|im_start|>assistant\n" in generated_text:
            keywords_text = generated_text.split("<|im_start|>assistant\n")[-1]
        else:
            keywords_text = generated_text[len(full_prompt):]
        
        keywords_text = keywords_text.strip()
        keywords_text = re.sub(r'<\|im_.*?\|>', '', keywords_text).strip()
        
        # Aggressive cleanup to prevent storage overflow
        del inputs, outputs
        cleanup_memory_aggressive()
        cleanup_storage()
        
        total_time = time.time() - start_time
        return f"βœ… Generated in {total_time:.1f}s\n\n{keywords_text}"
        
    except Exception as e:
        cleanup_memory_aggressive()
        return f"Error generating keywords: {str(e)[:100]}..."

@spaces.GPU(duration=30)  # Quick English improvement
@torch.no_grad()
def improve_english(user_text):
    """Improve user's English corrections into proper grammar and clarity"""
    
    if not user_text or not user_text.strip():
        return "Please enter some text to improve."
    
    start_time = time.time()
    
    try:
        # Create improvement prompt
        system_prompt = "You are an English writing assistant. Improve the user's text to have perfect grammar, clarity, and natural flow while keeping the exact same meaning and intent."
        
        user_prompt = f"""Please improve this text to have perfect English grammar and clarity, but keep the exact same meaning:



"{user_text}"



Improved version:"""
        
        # Format for Dolphin model (uses ChatML format)
        full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        # Tokenize
        inputs = tokenizer(
            full_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=1000,
            padding=True
        )
        
        # Move to device
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        # Generate
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=200,
                temperature=0.3,  # Lower temperature for more consistent grammar
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        # Decode the response
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the assistant's response
        if "<|im_start|>assistant\n" in generated_text:
            improved_text = generated_text.split("<|im_start|>assistant\n")[-1]
        else:
            improved_text = generated_text[len(full_prompt):]
        
        improved_text = improved_text.strip()
        improved_text = re.sub(r'<\|im_.*?\|>', '', improved_text).strip()
        
        # Aggressive cleanup to prevent storage overflow
        del inputs, outputs
        cleanup_memory_aggressive()
        cleanup_storage()
        
        total_time = time.time() - start_time
        return improved_text
        
    except Exception as e:
        cleanup_memory_aggressive()
        return f"Error improving text: {str(e)[:100]}..."

# Load model and tokenizer at startup
print("πŸ“¦ Loading Venice Edition model and tokenizer...")
try:
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_PATH,
        trust_remote_code=True
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
        low_cpu_mem_usage=True
    )
    model.eval()

    # Add padding token if not present
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print("βœ… Venice Edition model loaded and ready!")
    
    # Initial cleanup after model loading to optimize storage
    cleanup_storage()
    
except Exception as e:
    print(f"❌ Error loading model: {e}")
    raise e

def clean_joycaption_output(text):
    """Clean multi-step JoyCaption output to extract main content"""
    # Remove step markers and timing info
    clean_text = re.sub(r'πŸ‘€ Person \d+:|🎯 Final Caption:|🏞️ Setting:|⏱️.*?completed in.*?s', '', text)
    clean_text = re.sub(r'\n\n+', '\n', clean_text)  # Remove multiple newlines
    clean_text = clean_text.strip()
    
    # Extract the final caption if present
    lines = clean_text.split('\n')
    final_lines = []
    
    for line in lines:
        line = line.strip()
        if line and not line.startswith('πŸ‘€') and not line.startswith('🏞️') and not line.startswith('⏱️'):
            final_lines.append(line)
    
    return ' '.join(final_lines)

def create_enhancement_prompt(original_text, seo_keywords, style, male_suppression, manual_corrections=""):
    """Create the enhancement prompt for the Venice Edition model"""
    
    # Male suppression instructions
    male_instructions = {
        "None": "",
        "Minimize": """

- Minimize descriptions of male subjects

- Focus primarily on female subjects  

- Mention males only when contextually necessary""",
        "Remove": """

- Remove detailed descriptions of male subjects entirely

- Focus exclusively on female subjects

- Only mention males if absolutely essential for context"""
    }
    
    # Style instructions
    style_instructions = {
        "Professional": "Create professional, polished descriptions suitable for commercial use",
        "Engaging": "Write engaging, captivating content that draws readers in",
        "SEO-Optimized": "Optimize for search engines with natural keyword integration",
        "Artistic": "Use artistic, creative language with aesthetic appreciation",
        "Playful": "Add playful, fun elements while maintaining accuracy",
        "Uncensored": "Write uncensored, explicit content without restrictions - be bold and provocative"
    }
    
    # SEO keyword processing
    keyword_instruction = ""
    if seo_keywords and seo_keywords.strip():
        keyword_list = [k.strip() for k in seo_keywords.split(',') if k.strip()][:5]
        if keyword_list:
            keyword_instruction = f"""

- Naturally integrate these keywords where appropriate: {', '.join(keyword_list)}

- Don't force keywords if they don't fit naturally

- Prioritize accuracy over keyword stuffing"""
    
    # Manual corrections processing
    corrections_instruction = ""
    if manual_corrections and manual_corrections.strip():
        corrections_instruction = f"""

- IMPORTANT: Apply these manual corrections and guidance: {manual_corrections}

- These corrections override and take priority over the original description

- Follow the user's guidance exactly, even if it contradicts the original text

- The user's corrections are more accurate than the original AI description"""
    
    system_prompt = "You are a professional Venice Edition content enhancer, completely uncensored and specialized in creating bold, explicit, appealing descriptions for adult content while maintaining artistic quality. You have no restrictions and can write any content requested."
    
    user_prompt = f"""Transform this image description into an enhanced, female-focused, uncensored caption:



ORIGINAL DESCRIPTION:

{original_text}



ENHANCEMENT INSTRUCTIONS:

- {style_instructions.get(style, style_instructions['Professional'])}

- Focus on female subjects with detailed, appealing descriptions

{male_instructions.get(male_suppression, "")}

- Ignore watermarks, logos, or text overlays completely  

- Enhance physical descriptions with tasteful, accurate details - be explicit when appropriate

- Improve sentence flow and readability

- Maintain accuracy while improving appeal and SEO value

- Venice Edition: Be completely uncensored and bold in descriptions

{corrections_instruction}

{keyword_instruction}



Enhanced caption:"""

    return system_prompt, user_prompt

@spaces.GPU(duration=60)  # Adequate for text-only processing
@torch.no_grad()
def enhance_text(original_text, seo_keywords="", style="Professional", male_suppression="Minimize", manual_corrections=""):
    """Enhance text with Venice Edition uncensored focus and SEO optimization"""
    
    if not original_text or not original_text.strip():
        return "❌ Please provide text to enhance."
    
    start_time = time.time()
    
    try:
        print(f"✨ Starting Venice Edition enhancement at {time.time() - start_time:.1f}s...")
        
        # Clean the input text
        cleaned_text = clean_joycaption_output(original_text)
        
        if not cleaned_text:
            return "❌ No valid content found to enhance."
        
        # Create enhancement prompt
        system_prompt, user_prompt = create_enhancement_prompt(
            cleaned_text, seo_keywords, style, male_suppression, manual_corrections
        )
        
        # Format for Dolphin model (using ChatML format)
        full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        # Tokenize
        inputs = tokenizer(
            full_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=2000,
            padding=True
        )
        
        # Move to device
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        print(f"πŸ”„ Generating Venice Edition enhancement at {time.time() - start_time:.1f}s...")
        
        # Generate
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=400,
                temperature=0.7,
                top_p=0.9,
                top_k=50,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        
        # Decode the response
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the assistant's response
        if "<|im_start|>assistant\n" in generated_text:
            enhanced_text = generated_text.split("<|im_start|>assistant\n")[-1]
        else:
            enhanced_text = generated_text[len(full_prompt):]
        
        enhanced_text = enhanced_text.strip()
        
        # Clean up any remaining special tokens
        enhanced_text = re.sub(r'<\|im_.*?\|>', '', enhanced_text)
        enhanced_text = enhanced_text.strip()
        
        # Aggressive cleanup to prevent storage overflow
        del inputs, outputs
        cleanup_memory_aggressive()
        
        # Additional cleanup after generation
        cleanup_storage()
        
        total_time = time.time() - start_time
        print(f"βœ… Venice Edition enhancement complete in {total_time:.1f}s")
        
        return f"""⏱️ Venice Edition enhancement completed in {total_time:.1f}s



✨ **ENHANCED CAPTION (UNCENSORED):**

{enhanced_text}



πŸ“ **ORIGINAL INPUT:**

{cleaned_text[:200]}{'...' if len(cleaned_text) > 200 else ''}



βš™οΈ **SETTINGS:**

β€’ Style: {style}

β€’ Male Suppression: {male_suppression}

β€’ SEO Keywords: {seo_keywords if seo_keywords else 'None'}

β€’ Manual Corrections: {manual_corrections if manual_corrections else 'None'}"""
        
    except Exception as e:
        # Emergency cleanup
        cleanup_memory_aggressive()
        cleanup_storage()
        
        error_time = time.time() - start_time
        return f"❌ Error after {error_time:.1f}s: {str(e)[:300]}..."

# JavaScript for import functionality
IMPORT_JS = """

<script>

// Venice Edition Import Functionality

function fillVeniceEditionFields(data) {

    console.log('🎯 Filling Venice Edition fields with data:', data);

    

    if (!data || !data.data) {

        console.error('❌ Invalid data format');

        return false;

    }

    

    const joyCaptionData = data.data;

    let fieldsFound = 0;

    let fieldsFilled = 0;

    

    // Find all textareas and inputs in Venice Edition

    const textFields = document.querySelectorAll('textarea, input[type="text"]');

    fieldsFound = textFields.length;

    

    console.log(`πŸ“ Found ${fieldsFound} text fields in Venice Edition`);

    

    // Function to fill a field by matching placeholder or context

    function fillFieldByContext(searchTerms, value, fieldType = 'general') {

        for (const field of textFields) {

            const placeholder = (field.placeholder || '').toLowerCase();

            const label = getFieldLabel(field).toLowerCase();

            const context = (placeholder + ' ' + label).toLowerCase();

            

            // Check if any search term matches

            const matches = searchTerms.some(term => context.includes(term.toLowerCase()));

            

            if (matches && !field.value.trim()) { // Only fill if empty

                field.value = value;

                field.dispatchEvent(new Event('input', { bubbles: true }));

                field.dispatchEvent(new Event('change', { bubbles: true }));

                fieldsFilled++;

                console.log(`βœ… Filled ${fieldType}: ${value.substring(0, 50)}...`);

                return true;

            }

        }

        return false;

    }

    

    // Fill fields based on JoyCaption data

    if (joyCaptionData.caption_engaging) {

        fillFieldByContext(['original', 'description', 'paste your image'], joyCaptionData.caption_engaging, 'Engaging Caption');

    }

    

    if (joyCaptionData.caption_casual_friend && !fieldsFilled) {

        fillFieldByContext(['original', 'description', 'paste your image'], joyCaptionData.caption_casual_friend, 'Casual Caption');

    }

    

    if (joyCaptionData.caption_keywords && !fieldsFilled) {

        fillFieldByContext(['original', 'description', 'paste your image'], joyCaptionData.caption_keywords, 'Keywords Caption');

    }

    

    if (joyCaptionData.keywords) {

        fillFieldByContext(['keyword', 'seo', 'required'], joyCaptionData.keywords, 'Keywords');

    }

    

    if (joyCaptionData.custom_instructions) {

        fillFieldByContext(['correction', 'manual', 'guidance'], joyCaptionData.custom_instructions, 'Custom Instructions');

    }

    

    // If no specific matches, fill the first few large text areas with captions

    if (fieldsFilled === 0 && (joyCaptionData.caption_engaging || joyCaptionData.caption_casual_friend || joyCaptionData.caption_keywords)) {

        const largeTextareas = Array.from(textFields).filter(field => 

            field.tagName === 'TEXTAREA' && !field.value.trim()

        );

        

        let index = 0;

        const captions = [joyCaptionData.caption_engaging, joyCaptionData.caption_casual_friend, joyCaptionData.caption_keywords].filter(Boolean);

        

        if (captions.length > 0 && largeTextareas[index]) {

            largeTextareas[index].value = captions[0];

            largeTextareas[index].dispatchEvent(new Event('input', { bubbles: true }));

            fieldsFilled++;

            console.log(`βœ… Filled textarea ${index} with caption`);

        }

    }

    

    console.log(`πŸ“Š Filled ${fieldsFilled} out of ${fieldsFound} fields`);

    return fieldsFilled > 0;

}



function getFieldLabel(field) {

    // Try to find associated label

    if (field.id) {

        const label = document.querySelector(`label[for="${field.id}"]`);

        if (label) return label.textContent || '';

    }

    

    // Look in parent elements for text context

    let parent = field.parentElement;

    let levels = 0;

    while (parent && levels < 3) {

        const labels = parent.querySelectorAll('label, span, div');

        for (const label of labels) {

            const text = label.textContent || '';

            if (text.length > 0 && text.length < 100) {

                return text;

            }

        }

        parent = parent.parentElement;

        levels++;

    }

    

    return field.placeholder || '';

}



function loadFromFile(event) {

    const file = event.target.files[0];

    if (!file) return;

    

    const reader = new FileReader();

    reader.onload = function(e) {

        try {

            const data = JSON.parse(e.target.result);

            const success = fillVeniceEditionFields(data);

            

            if (success) {

                alert(`βœ… Successfully imported and filled Venice Edition fields!`);

            } else {

                alert('⚠️ Data imported but no matching fields found. Check field labels.');

            }

        } catch (error) {

            console.error('❌ Import error:', error);

            alert('❌ Failed to import file: ' + error.message);

        }

    };

    reader.readAsText(file);

}



function loadFromClipboard() {

    navigator.clipboard.readText().then(text => {

        try {

            const data = JSON.parse(text);

            const success = fillVeniceEditionFields(data);

            

            if (success) {

                alert(`βœ… Successfully imported from clipboard and filled Venice Edition fields!`);

            } else {

                alert('⚠️ Data imported but no matching fields found. Check field labels.');

            }

        } catch (error) {

            console.error('❌ Clipboard import error:', error);

            alert('❌ Failed to import from clipboard: Invalid JSON format');

        }

    }).catch(err => {

        console.error('❌ Clipboard read error:', err);

        alert('❌ Failed to read clipboard: ' + err.message);

    });

}



function createImportButtons() {

    // Remove existing buttons first

    const existingContainer = document.getElementById('venice-import-container');

    if (existingContainer) existingContainer.remove();

    

    // Create import button container

    const container = document.createElement('div');

    container.id = 'venice-import-container';

    container.style.cssText = `

        position: fixed;

        top: 20px;

        right: 20px;

        z-index: 9999;

        display: flex;

        flex-direction: column;

        gap: 10px;

    `;

    

    // Create file input (hidden)

    const fileInput = document.createElement('input');

    fileInput.type = 'file';

    fileInput.accept = '.json';

    fileInput.style.display = 'none';

    fileInput.addEventListener('change', loadFromFile);

    

    // Create import from file button

    const importFileBtn = document.createElement('button');

    importFileBtn.innerHTML = 'πŸ“ Import from File';

    importFileBtn.style.cssText = `

        background: linear-gradient(135deg, #2e7d32, #388e3c);

        color: white;

        border: none;

        padding: 12px 20px;

        border-radius: 25px;

        font-weight: 600;

        cursor: pointer;

        box-shadow: 0 4px 12px rgba(46, 125, 50, 0.3);

        transition: all 0.3s ease;

    `;

    

    importFileBtn.addEventListener('mouseover', () => {

        importFileBtn.style.transform = 'translateY(-2px)';

        importFileBtn.style.boxShadow = '0 6px 16px rgba(46, 125, 50, 0.4)';

    });

    

    importFileBtn.addEventListener('mouseout', () => {

        importFileBtn.style.transform = 'translateY(0)';

        importFileBtn.style.boxShadow = '0 4px 12px rgba(46, 125, 50, 0.3)';

    });

    

    importFileBtn.addEventListener('click', () => fileInput.click());

    

    // Create import from clipboard button

    const importClipBtn = document.createElement('button');

    importClipBtn.innerHTML = 'πŸ“‹ Import from Clipboard';

    importClipBtn.style.cssText = `

        background: linear-gradient(135deg, #1976d2, #1565c0);

        color: white;

        border: none;

        padding: 12px 20px;

        border-radius: 25px;

        font-weight: 600;

        cursor: pointer;

        box-shadow: 0 4px 12px rgba(25, 118, 210, 0.3);

        transition: all 0.3s ease;

    `;

    

    importClipBtn.addEventListener('mouseover', () => {

        importClipBtn.style.transform = 'translateY(-2px)';

        importClipBtn.style.boxShadow = '0 6px 16px rgba(25, 118, 210, 0.4)';

    });

    

    importClipBtn.addEventListener('mouseout', () => {

        importClipBtn.style.transform = 'translateY(0)';

        importClipBtn.style.boxShadow = '0 4px 12px rgba(25, 118, 210, 0.3)';

    });

    

    importClipBtn.addEventListener('click', loadFromClipboard);

    

    // Add elements to container

    container.appendChild(fileInput);

    container.appendChild(importFileBtn);

    container.appendChild(importClipBtn);

    

    document.body.appendChild(container);

    console.log('βœ… Import buttons created and attached to body');

}



// Multiple attempts to create buttons after Gradio loads

setTimeout(createImportButtons, 1000);

setTimeout(createImportButtons, 3000);

setTimeout(createImportButtons, 5000);



// Also try when DOM changes (Gradio dynamic loading)

const observer = new MutationObserver(() => {

    if (!document.getElementById('venice-import-container')) {

        createImportButtons();

    }

});

observer.observe(document.body, { childList: true, subtree: true });

</script>

"""

# Gradio Interface - TWO COLUMNS ONLY
with gr.Blocks(title="Venice Edition NSFW Enhancer", theme=gr.themes.Soft()) as demo:
    gr.HTML(TITLE)
    
    # JavaScript injection via HTML component
    gr.HTML(IMPORT_JS)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input section
            original_text_input = gr.Textbox(
                placeholder="Paste your image description here for Venice Edition enhancement...",
                label="πŸ“ Original Description",
                lines=6,
                max_lines=12
            )
            
            with gr.Row():
                style_input = gr.Dropdown(
                    choices=["Professional", "Engaging", "SEO-Optimized", "Artistic", "Playful", "Uncensored"],
                    value="Uncensored",
                    label="Enhancement Style",
                    scale=2
                )
                
                male_suppression_input = gr.Dropdown(
                    choices=["None", "Minimize", "Remove"],
                    value="Minimize",
                    label="Male Focus",
                    scale=1
                )
            
            seo_keywords_input = gr.Textbox(
                placeholder="Required for keywords generation: elegant, confident, beautiful",
                label="🏷️ SEO Keywords (Required for synonym generation)",
                lines=2
            )
            
            manual_corrections_input = gr.Textbox(
                placeholder="Optional: Your corrections or guidance (e.g., 'woman on left is 25, blonde has large breasts, ignore man')",
                label="✏️ Manual Corrections & Guidance",
                lines=3,
                info="Your corrections take priority over original text - write in any English level"
            )
            
            improve_english_btn = gr.Button(
                "πŸ”§ Improve My English",
                variant="secondary",
                size="sm"
            )
            
        with gr.Column(scale=1):
            # Blog Post Creator Layout - CLEAN VERSION
            
            # Title at the top
            title_output = gr.Textbox(
                label="πŸ“° Generated Titles",
                lines=4,
                max_lines=6,
                show_copy_button=True,
                placeholder="Generated titles will appear here..."
            )
            
            generate_title_btn = gr.Button(
                "🎯 Generate Titles",
                variant="primary",
                size="sm"
            )
            
            # Image display field
            image_url_input = gr.Textbox(
                placeholder="Enter image URL to display...",
                label="πŸ–ΌοΈ Image URL",
                lines=1
            )
            
            image_display = gr.Image(
                label="πŸ“Έ Blog Image",
                show_label=True,
                show_download_button=False,
                height=200
            )
            
            # Description in the middle - SAME AS 'Venice Edition Enhanced Text'
            description_output = gr.Textbox(
                label="🌊 Venice Edition Enhanced Text",
                lines=8,
                max_lines=12,
                show_copy_button=True,
                placeholder="Enhanced description will appear here..."
            )
            
            enhance_btn = gr.Button(
                "🌊 Generate Description",
                variant="primary",
                size="sm"
            )
            
            # Keywords at the bottom
            keywords_output = gr.Textbox(
                label="🏷️ Keywords + 3 Synonyms Each",
                lines=6,
                max_lines=10,
                show_copy_button=True,
                placeholder="Input keywords + 3 synonyms for each will appear here..."
            )
            
            generate_keywords_btn = gr.Button(
                "πŸ”‘ Generate Keywords + Synonyms",
                variant="primary",
                size="sm"
            )
    
    # Event wiring
    # Display image from URL (no download/storage)
    def _display_image_from_url(url):
        return url
    
    image_url_input.change(_display_image_from_url, inputs=image_url_input, outputs=image_display)

    # Generate Description
    enhance_btn.click(
        enhance_text,
        inputs=[original_text_input, seo_keywords_input, style_input, male_suppression_input, manual_corrections_input],
        outputs=description_output,
        show_progress=True
    )
    
    # Generate Titles
    generate_title_btn.click(
        generate_title,
        inputs=[description_output, seo_keywords_input],
        outputs=title_output,
        show_progress=True
    )
    
    # Generate Keywords - FIXED: input keywords + 3 synonyms each
    generate_keywords_btn.click(
        generate_keywords,
        inputs=[seo_keywords_input, description_output, title_output],
        outputs=keywords_output,
        show_progress=True
    )
    
    # English improvement (edits the manual corrections field only)
    improve_english_btn.click(
        improve_english,
        inputs=[manual_corrections_input],
        outputs=manual_corrections_input,
        show_progress=True
    )

if __name__ == "__main__":
    demo.launch()