Update app.py

app.py

@@ -1,15 +1,15 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# Importar bibliotecas
 import torch
-from clip_interrogator import Config, Interrogator
-import random
 import re
+import random
 import requests
 import shutil
+from clip_interrogator import Config, Interrogator
+from transformers import pipeline, set_seed, AutoTokenizer, AutoModelForSeq2SeqLM
 from PIL import Image
+import gradio as gr
 
-
-# Definir la función para generar prompt desde imagen
+# Configurar CLIP
 config = Config()
 config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 config.blip_offload = False if torch.cuda.is_available() else True
@@ -19,6 +19,7 @@ config.blip_num_beams = 64
 config.clip_model_name = "ViT-H-14/laion2b_s32b_b79k"
 ci = Interrogator(config)
 
+# Función para generar prompt desde imagen
 def get_prompt_from_image(image, mode):
     image = image.convert('RGB')
     if mode == 'best':
@@ -31,29 +32,18 @@ def get_prompt_from_image(image, mode):
         prompt = ci.interrogate_negative(image)
     return prompt
 
-
-# Definir la función para generar prompt desde texto
-model = AutoModelForSeq2SeqLM.from_pretrained('Helsinki-NLP/opus-mt-zh-en').eval()
-tokenizer = AutoTokenizer.from_pretrained('Helsinki-NLP/opus-mt-zh-en')
-
-def translate(text):
-    with torch.no_grad():
-        encoded = tokenizer([text], return_tensors='pt')
-        sequences = model.generate(**encoded)
-        return tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
-
+# Función para generar texto
 text_pipe = pipeline('text-generation', model='succinctly/text2image-prompt-generator')
 
 def text_generate(input):
     seed = random.randint(100, 1000000)
     set_seed(seed)
-    text_in_english = translate(input)
     for count in range(6):    
-        sequences = text_pipe(text_in_english, max_length=random.randint(60, 90), num_return_sequences=8)
+        sequences = text_pipe(input, max_length=random.randint(60, 90), num_return_sequences=8)
         list = []
         for sequence in sequences:
             line = sequence['generated_text'].strip()
-            if line != text_in_english and len(line) > (len(text_in_english) + 4) and line.endswith((':', '-', '—')) is False:
+            if line != input and len(line) > (len(input) + 4) and line.endswith((':', '-', '—')) is False:
                 list.append(line)
 
         result = "\n".join(list)
@@ -64,25 +54,24 @@ def text_generate(input):
         if count == 5:
             return result
 
+# Crear interfaz gradio
+with gr.Blocks() as block:
+    with gr.Column():
+        gr.HTML('<h1>MidJourney / SD2 Helper Tool</h1>')
+        with gr.Tab('Generate from Image'):
+            with gr.Row():
+                input_image = gr.Image(type='pil')
+                with gr.Column():
+                    input_mode = gr.Radio(['best', 'fast', 'classic', 'negative'], value='best', label='Mode')
+            img_btn = gr.Button('Discover Image Prompt')
+            output_image = gr.Textbox(lines=6, label='Generated Prompt')
 
-# Definir la función que permite al usuario cargar una imagen desde una URL
-def load_image_from_url(url):
-    response = requests.get(url, stream=True)
-    if response.status_code == 200:
-        with open('./image.jpg', 'wb') as f:
-            response.raw.decode_content = True
-            shutil.copyfileobj(response.raw, f) 
-        return Image.open('./image.jpg')
-    else:
-        raise ValueError("No se pudo cargar la imagen")
+        with gr.Tab('Generate from Text'):
+            input_text = gr.Textbox(lines=6, label='Your Idea', placeholder='Enter your content here...')
+            output_text = gr.Textbox(lines=6, label='Generated Prompt')
+            text_btn = gr.Button('Generate Prompt')
 
+    img_btn.click(fn=get_prompt_from_image, inputs=[input_image, input_mode], outputs=output_image)
+    text_btn.click(fn=text_generate, inputs=input_text, outputs=output_text)
 
-# Crear la interfaz de usuario de Gradio
-with gr.Interface(
-    [get_prompt_from_image, text_generate],
-    [
-        gr.inputs.Image(type='pil', label='Imagen'),
-        gr.inputs.Radio(['best', 'fast', 'classic', 'negative'], value='best', label='Modo'),
-        gr.inputs.Textbox(lines=6, label='Texto de entrada'),
-    ],
-    [
+block.queue(max_size=64).launch(show_api=False, enable_queue=True, debug=True, share=True, server_name='0.0.0.0')