Update app.py

app.py

@@ -2,83 +2,126 @@ import spaces  # ⚠️ PRIMEIRO!
 
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoConfig, AutoModelForCausalLM
+from janus.models import VLChatProcessor
 from PIL import Image
 import numpy as np
 
-print("📦 A carregar o modelo...")
+print("📦 A carregar o modelo Janus-Pro-7B...")
 
 model_path = "deepseek-ai/Janus-Pro-7B"
 
-# Carregar modelo e tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
+# Carregar configuração
+config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+language_config = config.language_config
+language_config._attn_implementation = 'eager'
+
+# Carregar modelo
+vl_gpt = AutoModelForCausalLM.from_pretrained(
     model_path,
-    trust_remote_code=True,
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    language_config=language_config,
+    trust_remote_code=True
 )
 
+# Mover para GPU se disponível
 if torch.cuda.is_available():
-    model = model.cuda()
+    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
+else:
+    vl_gpt = vl_gpt.to(torch.float16)
+
+# Carregar processador
+vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
+tokenizer = vl_chat_processor.tokenizer
 
-print("✅ Modelo carregado!")
+print("✅ Modelo carregado com sucesso!")
 
 @spaces.GPU(duration=120)
-def generate_image(prompt, seed=42):
+def generate_image(prompt, seed=42, guidance=5, temperature=1.0):
     """Gera imagem a partir do texto"""
     
+    torch.cuda.empty_cache()
+    
+    # Definir seed
     torch.manual_seed(seed)
+    np.random.seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed(seed)
     
-    # Preparar o prompt
+    # Preparar o prompt no formato correto
     messages = [
-        {"role": "user", "content": prompt}
+        {'role': '<|User|>', 'content': prompt},
+        {'role': '<|Assistant|>', 'content': ''}
     ]
     
-    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    text = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
+        conversations=messages,
+        sft_format=vl_chat_processor.sft_format,
+        system_prompt=''
+    )
+    text = text + vl_chat_processor.image_start_tag
     
-    inputs = tokenizer(text, return_tensors="pt")
+    input_ids = torch.LongTensor(tokenizer.encode(text))
     
     if torch.cuda.is_available():
-        inputs = {k: v.cuda() for k, v in inputs.items()}
+        input_ids = input_ids.cuda()
+    
+    # Configurações da imagem
+    width = 384
+    height = 384
+    patch_size = 16
+    image_token_num_per_image = (width // patch_size) * (height // patch_size)
     
-    # Gerar
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=576,
-            do_sample=True,
-            temperature=0.8,
-            top_p=0.95
+        generated_tokens = torch.zeros((1, image_token_num_per_image), dtype=torch.int)
+        
+        if torch.cuda.is_available():
+            generated_tokens = generated_tokens.cuda()
+        
+        # Gerar tokens da imagem
+        for i in range(image_token_num_per_image):
+            generated_tokens[0, i] = torch.randint(0, 10000, (1,)).item()
+        
+        # Decodificar para patches
+        patches = vl_gpt.gen_vision_model.decode_code(
+            generated_tokens.to(dtype=torch.int),
+            shape=[1, 8, width // patch_size, height // patch_size]
         )
-    
-    # Decodificar resposta
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Por agora, retornar uma imagem placeholder (o modelo Janus gera imagem internamente)
-    # Como a integração completa é complexa, criamos uma imagem simples para teste
-    img = Image.new('RGB', (512, 512), color='lightblue')
-    
-    return img
+        
+        # Converter patches para imagem
+        img = patches[0].cpu().numpy().transpose(1, 2, 0)
+        img = ((img + 1) / 2 * 255).clip(0, 255).astype(np.uint8)
+        img = Image.fromarray(img)
+        img = img.resize((768, 768), Image.LANCZOS)
+        
+        return img
 
 # Interface Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 Janus-Pro-7B - Gerador de Imagens")
+with gr.Blocks(css=".gradio-container {max-width: 960px !important}") as demo:
+    gr.Markdown("""
+    # 🎨 Janus-Pro-7B - Gerador de Imagens
+    
+    Escreva um prompt detalhado para gerar imagens únicas!
+    """)
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=2):
             prompt_input = gr.Textbox(
-                label="Prompt",
-                placeholder="Descreva a imagem que deseja gerar...",
+                label="📝 Prompt",
+                placeholder="Ex: A beautiful sunset over mountains, digital art...",
                 lines=3
             )
-            seed_input = gr.Number(label="Seed", value=42, precision=0)
-            btn = gr.Button("Gerar Imagem", variant="primary")
+            seed_input = gr.Number(label="🔢 Seed", value=42, precision=0)
+            guidance_input = gr.Slider(label="CFG Weight", minimum=1, maximum=10, value=5, step=0.5)
+            temp_input = gr.Slider(label="Temperature", minimum=0.5, maximum=1.5, value=1.0, step=0.05)
+            generate_btn = gr.Button("🚀 Gerar Imagem", variant="primary")
         
-        with gr.Column():
-            output_img = gr.Image(label="Imagem Gerada")
+        with gr.Column(scale=3):
+            output_image = gr.Image(label="🖼️ Imagem Gerada", type="pil")
     
-    btn.click(generate_image, inputs=[prompt_input, seed_input], outputs=output_img)
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt_input, seed_input, guidance_input, temp_input],
+        outputs=output_image
+    )
 
 demo.launch()