Update app.py

app.py

@@ -1,19 +1,29 @@
 import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Carrega o modelo FLUX.1 diretamente
+# Função para carregar o modelo base com LoRa
+def load_model_with_lora(base_model_path, lora_model_path):
+    # Carrega o modelo base
+    pipeline = StableDiffusionPipeline.from_pretrained(base_model_path, torch_dtype=torch.float16).to("cuda")
+    
+    # Carrega a LoRa e aplica ao modelo base
+    pipeline.load_lora_weights(lora_model_path)
+    
+    return pipeline
+
+# Função de inferência de imagem com parâmetros ajustáveis
 def infer_image(prompt, steps, cfg_scale, seed, width, height):
-    # Aqui, você utiliza o modelo do Hugging Face carregado
-    result = gr.load("models/rorito/testSCG-Anatomy-Flux1")(
-        prompt=prompt,
-        steps=steps,
-        cfg_scale=cfg_scale,
-        seed=seed,
-        width=width,
-        height=height
-    )
-    return result
+    pipeline = load_model_with_lora("black-forest-labs/FLUX.1-dev", "rorito/testSCG-Anatomy-Flux1")
+    
+    # Configurações adicionais
+    generator = torch.manual_seed(seed)
+    result = pipeline(prompt, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator)
+    
+    return result.images[0]
 
-# Interface com Gradio, permitindo ajustes nas configurações
+# Interface do Gradio com parâmetros ajustáveis
 interface = gr.Interface(
     fn=infer_image,
     inputs=[
@@ -27,5 +37,5 @@ interface = gr.Interface(
     outputs="image"
 )
 
-# Lança a aplicação Gradio
+# Lançar o aplicativo Gradio
 interface.launch()