Update app.py

app.py

@@ -1,62 +1,24 @@
 import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from PIL import Image
+from diffusers import StableDiffusionPipeline
 import torch
-import os
 
-# 1. Definição do Nome do Modelo (NOVO MODELO PEQUENO)
-# A versão 'small' printed é muito mais leve e confiável para a maioria dos Spaces.
-MODEL_NAME = "microsoft/trocr-small-printed" 
+# 1. Carregar o Pipeline (use um modelo adequado ao seu hardware)
+pipeline = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", 
+    torch_dtype=torch.float16
+).to("cuda")
 
-# 2. Carregamento do Processador e do Modelo (Feito Apenas Uma Vez)
-try:
-    processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
-    model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)
-    print("Modelo TrOCR SMALL carregado com sucesso.")
-except Exception as e:
-    # Se ainda falhar, o erro exato estará nos Logs.
-    print(f"Erro crítico ao carregar o modelo: {e}") 
-    processor = None
-    model = None
+def generate_image(prompt):
+    # 2. Gerar a Imagem
+    image = pipeline(prompt).images[0]
+    return image
 
-# Verifica se a GPU está disponível e define o dispositivo
-# A versão 'small' é perfeitamente funcional na CPU, o que também economiza VRAM.
-device = "cuda" if torch.cuda.is_available() else "cpu"
-if model:
-    model.to(device)
-    print(f"Modelo movido para o dispositivo: {device}")
-
-
-# 3. Função de Inferência (OCR)
-def realizar_ocr(imagem: Image.Image) -> str:
-    # ... (O restante da função de inferência permanece o mesmo)
-    if not model or not processor:
-        return "Erro: O modelo não foi carregado corretamente na inicialização."
-
-    if imagem is None:
-        return "Nenhuma imagem fornecida."
-
-    # Processamento
-    pixel_values = processor(images=imagem, return_tensors="pt").pixel_values
-    pixel_values = pixel_values.to(device)
-
-    # Inferência
-    generated_ids = model.generate(pixel_values)
-    
-    # Decodificação
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-    return generated_text
-
-# 4. Configuração da Interface Gradio
+# 3. Criar a Interface Gradio
 iface = gr.Interface(
-    fn=realizar_ocr,
-    inputs=gr.Image(type="pil", label="Faça o Upload da Imagem para OCR"),
-    outputs=gr.Textbox(label="Texto Reconhecido (OCR)"),
-    title="TrOCR Pequeno - Reconhecimento de Texto Impresso",
-    description="Utiliza a versão SMALL do TrOCR para extrair texto de imagens de documentos impressos (baixo consumo de memória)."
+    fn=generate_image, 
+    inputs="text", 
+    outputs="image",
+    title="Meu Gerador de Imagens com Stable Diffusion"
 )
 
-# Inicia o aplicativo Gradio
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()