ivanoctaviogaitansantos commited on
Commit
8c0cd05
verified
1 Parent(s): cc63205

Create app py

Browse files
Files changed (1) hide show
  1. app.py +116 -0
app.py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import gc
4
+ from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
5
+ from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection
6
+ import logging
7
+
8
+ logging.basicConfig(level=logging.INFO)
9
+ logger = logging.getLogger(__name__)
10
+
11
+ class AIImageGeneratorNSFW:
12
+ def __init__(self):
13
+ self.pipeline = None
14
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
15
+ # 隆Aqu铆 va la l铆nea!
16
+ self.model_id = "segmind/Segmind-DE-XL"
17
+ self.is_model_loaded = False
18
+ logger.info(f"Inicializando en dispositivo: {self.device}")
19
+
20
+ def load_model(self):
21
+ if self.is_model_loaded:
22
+ return True
23
+ try:
24
+ logger.info("Cargando modelo NSFW...")
25
+ torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
26
+
27
+ tokenizer_1 = CLIPTokenizer.from_pretrained(self.model_id, subfolder="tokenizer", use_fast=False)
28
+ tokenizer_2 = CLIPTokenizer.from_pretrained(self.model_id, subfolder="tokenizer_2", use_fast=False)
29
+
30
+ text_encoder_1 = CLIPTextModel.from_pretrained(self.model_id, subfolder="text_encoder", torch_dtype=torch_dtype, low_cpu_mem_usage=True)
31
+ text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(self.model_id, subfolder="text_encoder_2", torch_dtype=torch_dtype, low_cpu_mem_usage=True)
32
+
33
+ self.pipeline = StableDiffusionXLPipeline.from_pretrained(
34
+ self.model_id,
35
+ tokenizer=[tokenizer_1, tokenizer_2],
36
+ text_encoder=[text_encoder_1, text_encoder_2],
37
+ torch_dtype=torch_dtype,
38
+ scheduler=EulerDiscreteScheduler.from_pretrained(self.model_id, subfolder="scheduler"),
39
+ safety_checker=None,
40
+ use_safetensors=True,
41
+ variant="fp16" if self.device == "cuda" else None
42
+ )
43
+ self.pipeline.to(self.device)
44
+ self.is_model_loaded = True
45
+ logger.info("Modelo NSFW cargado correctamente.")
46
+ return True
47
+ except Exception as e:
48
+ logger.error(f"Error cargando modelo NSFW: {e}")
49
+ return False
50
+
51
+ def generate_image(self, prompt, width=1024, height=576, steps=35, guidance_scale=12.0):
52
+ if not self.is_model_loaded and not self.load_model():
53
+ return None
54
+ try:
55
+ with torch.inference_mode():
56
+ generator = torch.Generator(self.device).manual_seed(torch.randint(0, 2**32, (1,)).item())
57
+ result = self.pipeline(
58
+ prompt=prompt,
59
+ width=(width // 8) * 8,
60
+ height=(height // 8) * 8,
61
+ num_inference_steps=steps,
62
+ guidance_scale=guidance_scale,
63
+ generator=generator,
64
+ output_type="pil"
65
+ )
66
+ gc.collect()
67
+ return result.images[0]
68
+ except Exception as e:
69
+ logger.error(f"Error generando imagen NSFW: {e}")
70
+ gc.collect()
71
+ return None
72
+
73
+ # Instancia global
74
+ generator_nsfw = None
75
+
76
+ def initialize_generator_nsfw():
77
+ global generator_nsfw
78
+ if generator_nsfw is None:
79
+ generator_nsfw = AIImageGeneratorNSFW()
80
+ return generator_nsfw
81
+
82
+ def generate_image_nsfw(prompt, width, height, steps, guidance_scale):
83
+ gen = initialize_generator_nsfw()
84
+ if not prompt.strip():
85
+ return None
86
+ return gen.generate_image(
87
+ prompt=prompt,
88
+ width=int(width),
89
+ height=int(height),
90
+ steps=int(steps),
91
+ guidance_scale=float(guidance_scale)
92
+ )
93
+
94
+ def create_nsfw_interface():
95
+ with gr.Blocks(title="Generador de Im谩genes NSFW con IA - Stable Diffusion") as iface:
96
+ gr.Markdown("# 馃帹 Generador NSFW basado en Stable Diffusion\n_Uso responsable y solo para adultos_")
97
+
98
+ prompt = gr.Textbox(label="Prompt para la imagen NSFW", placeholder="Describe el contenido expl铆cito...", lines=3)
99
+ width = gr.Slider(512, 1536, value=1024, step=8, label="Ancho (pixeles)")
100
+ height = gr.Slider(512, 1536, value=576, step=8, label="Alto (pixeles)")
101
+ steps = gr.Slider(10, 50, value=35, step=1, label="Pasos de inferencia")
102
+ guidance_scale = gr.Slider(1.0, 20.0, value=12.0, step=0.1, label="Escala de gu铆a")
103
+
104
+ btn_generate = gr.Button("Generar Imagen NSFW")
105
+ img_output = gr.Image(label="Imagen generada")
106
+
107
+ btn_generate.click(
108
+ fn=generate_image_nsfw,
109
+ inputs=[prompt, width, height, steps, guidance_scale],
110
+ outputs=img_output
111
+ )
112
+ return iface
113
+
114
+ if __name__ == "__main__":
115
+ nsfw_app = create_nsfw_interface()
116
+ nsfw_app.launch()