Update app.py

app.py

@@ -1,201 +1,167 @@
+# app.py - Carrega base + UNet de repo privado separado
+# Data e hora atuais para referência: Sunday, May 4, 2025 at 8:23:22 PM -03
 import os, random, uuid, json
 import gradio as gr
 import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+# Importar UNet também
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerAncestralDiscreteScheduler
+import time
+from huggingface_hub import HfApi
 
-DESCRIPTION = None
+# --- Configurações ---
+base_model_id = "sd-community/sdxl-flash" # Ou o base que o Space usava
+# ID do Repositório PRIVADO que contém APENAS o UNet treinado
+tuned_unet_repo_id = "borsojj/unet" # <<< Repo com seu UNet
+
+DESCRIPTION = f"Interface usando base `{base_model_id}` com UNet de `{tuned_unet_repo_id}`."
 if not torch.cuda.is_available():
-    DESCRIPTION = "\nRunning on CPU 🥶 This demo may not work on CPU."
+    DESCRIPTION += "\n**Atenção:** Rodando em CPU 🥶 - A geração pode ser muito lenta ou falhar."
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "sd-community/sdxl-flash",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-    add_watermarker=False
-)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-
-if torch.cuda.is_available():
-    pipe.to("cuda")
+print(f"Carregando pipeline base de: {base_model_id}...")
+hf_token = os.getenv("HF_TOKEN") # Pega token dos segredos do Space
+if hf_token:
+    print("Segredo HF_TOKEN encontrado.")
 else:
-    pipe.to("cpu")
-
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
-
+    # AVISO IMPORTANTE se o repo UNET for privado!
+    print("AVISO: Segredo HF_TOKEN NÃO encontrado. O carregamento do UNet treinado falhará se o repositório for privado.")
+
+start_time = time.time()
+pipe = None
+loading_error_message = ""
+try:
+    # 1. Carrega o pipeline base completo
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        base_model_id,
+        torch_dtype=torch_dtype,
+        use_safetensors=True,
+        add_watermarker=False,
+        token=hf_token # Passa token caso o base precise também
+    )
+    print(f"Pipeline base '{base_model_id}' carregado.")
+
+    # 2. Tenta carregar e substituir o UNet do repo separado e PRIVADO
+    if hf_token: # Só tenta carregar se houver token
+        print(f"Tentando carregar UNet treinado do repo privado: {tuned_unet_repo_id}")
+        try:
+            # Carrega o UNet do repo ID, usando o token
+            tuned_unet = UNet2DConditionModel.from_pretrained(
+                tuned_unet_repo_id,
+                torch_dtype=torch_dtype, # Carrega com mesmo dtype
+                token=hf_token # Usa o token para acessar repo privado
+                # low_cpu_mem_usage=False # Pode precisar desativar se der erro OOM aqui
+            )
+            print("UNet treinado carregado. Substituindo UNet no pipeline...")
+            pipe.unet = tuned_unet # A SUBSTITUIÇÃO
+            print("UNet substituído com sucesso.")
+        except Exception as unet_load_e:
+            loading_error_message = f"**<font color='red'>ERRO:</font>** Falha ao carregar UNet de `{tuned_unet_repo_id}` (Verifique token e repo). Usando UNet base. Erro: `{unet_load_e}`"
+            print(loading_error_message)
+            DESCRIPTION += "\n" + loading_error_message
+            # Continua com o UNet base se falhar
+    else:
+         # Se não há token, não pode carregar UNet privado
+         loading_error_message = f"**<font color='orange'>AVISO:</font>** HF_TOKEN não encontrado. Não é possível carregar UNet do repositório privado `{tuned_unet_repo_id}`. Usando UNet base."
+         print(loading_error_message)
+         DESCRIPTION += "\n" + loading_error_message
+
+
+    # Configura o scheduler (como antes)
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    print(f"Scheduler configurado para: {pipe.scheduler.__class__.__name__}")
+
+    print(f"Movendo pipeline final para o device: {device}")
+    pipe.to(device) # Move para o device APÓS substituir o UNet
+    print("Pipeline pronto no device.")
+
+except Exception as e:
+    # Erro ao carregar o pipeline BASE
+    print(f"Erro CRÍTICO ao carregar o pipeline base de '{base_model_id}': {e}")
+    loading_error_message = f"**<font color='red'>ERRO CRÍTICO:</font>** Não foi possível carregar o pipeline base `{base_model_id}`. Erro: `{e}`."
+    DESCRIPTION += "\n" + loading_error_message
+    pipe = None
+
+
+# Função generate (sem alterações significativas)
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    if randomize_seed: seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=90)
 def generate(
     prompt: str,
     negative_prompt: str = "",
-    use_negative_prompt: bool = False,
+    use_negative_prompt: bool = True,
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
-    guidance_scale: float = 3,
+    guidance_scale: float = 7.0,
     num_inference_steps: int = 25,
     randomize_seed: bool = False,
-    use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
+    if pipe is None: raise gr.Error(f"Pipeline não carregado. {loading_error_message}")
     pipe.to(device)
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator().manual_seed(seed)   
-
-    options = {
-        "prompt":prompt,
-        "negative_prompt":negative_prompt,
-        "width":width,
-        "height":height,
-        "guidance_scale":guidance_scale,
-        "num_inference_steps":num_inference_steps,
-        "generator":generator,
-        "use_resolution_binning":use_resolution_binning,
-        "output_type":"pil",
-
-    }
-    
-    images = pipe(**options).images
-
-    image_paths = [save_image(img) for img in images]
-    return image_paths, seed
-
-
-examples = [
-    "a cat eating a piece of cheese",
-    "a ROBOT riding a BLUE horse on Mars, photorealistic",
-    "a cartoon of a IRONMAN fighting with HULK, wall painting",
-    "a cute robot artist painting on an easel, concept art",
-    "Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k",
-    "An alien grasping a sign board contain word 'Flash', futuristic, neonpunk, detailed",
-    "Kids going to school, Anime style"
-]
-
-css = '''
-.gradio-container {
-    max-width: 700px !important;
-    margin: 0 auto !important; 
-}
-h1{text-align:left}
-'''
+    generator = torch.Generator(device=device).manual_seed(seed)
+    if not use_negative_prompt: negative_prompt = None
+    options = {"prompt":prompt, "negative_prompt":negative_prompt, "width":width, "height":height, "guidance_scale":guidance_scale, "num_inference_steps":num_inference_steps, "generator":generator, "output_type":"pil"}
+    print(f"Gerando imagem com seed: {seed}, Steps: {num_inference_steps}, Guidance: {guidance_scale}")
+    start_gen_time = time.time()
+    try:
+        images = pipe(**options).images
+        print(f"Gerado {len(images)} imagem(s) em {time.time() - start_gen_time:.2f} segundos.")
+        return images, seed
+    except Exception as e:
+        print(f"Erro durante a geração: {e}")
+        raise gr.Error(f"Erro durante a geração: {e}")
+
+
+# Interface Gradio (sem alterações significativas)
+examples = [ "photo of a futuristic city skyline at sunset, high detail", "an oil painting of a cute cat wearing a wizard hat", "Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k", "An alien grasping a sign board contain word 'Flash', futuristic, neonpunk, detailed" ]
+css = ".gradio-container { max-width: 800px !important; margin: 0 auto !important; } h1{ text-align:center }"
 with gr.Blocks(css=css) as demo:
-    gr.Markdown(f"""# SDXL Flash
-        ### First Image processing takes time then images generate faster.
+    gr.Markdown(f"""# SDXL Base com UNet Fine-tuned (`{tuned_unet_repo_id}`)
+        Base: `{base_model_id}`
         {DESCRIPTION}""")
     with gr.Group():
         with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Run", scale=0)
-        result = gr.Gallery(label="Result", columns=1)
-    with gr.Accordion("Advanced options", open=False):
+            prompt = gr.Text(label="Prompt", show_label=False, max_lines=3, placeholder="Descreva a imagem...", container=False)
+            run_button = gr.Button("Gerar Imagem", variant="primary", scale=0)
+        result = gr.Gallery(label="Resultado", show_label=False, elem_id="gallery", columns=1, height=768)
+    with gr.Accordion("Opções Avançadas", open=False):
         with gr.Row():
-            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=5,
-                lines=4,
-                placeholder="Enter a negative prompt",
-                value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW",
-                visible=True,
-            )
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0,
-        )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row(visible=True):
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=64,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=64,
-                value=1024,
-            )
+            use_negative_prompt = gr.Checkbox(label="Usar prompt negativo", value=True)
+            negative_prompt = gr.Text(label="Prompt Negativo", max_lines=3, lines=2, placeholder="O que NÃO ver...", value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW", visible=True)
+        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+        randomize_seed = gr.Checkbox(label="Seed Aleatória", value=True)
         with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance Scale",
-                minimum=0.1,
-                maximum=6,
-                step=0.1,
-                value=3.0,
-            )
-            num_inference_steps = gr.Slider(
-                label="Number of inference steps",
-                minimum=1,
-                maximum=15,
-                step=1,
-                value=8,
-            )
-
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=[result, seed],
-        fn=generate,
-        cache_examples=True,
-    )
-
-    use_negative_prompt.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt,
-        outputs=negative_prompt,
-        api_name=False,
-    )
-
-    gr.on(
-        triggers=[
-            prompt.submit,
-            negative_prompt.submit,
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            randomize_seed,
-        ],
-        outputs=[result, seed],
-        api_name="run",
-    )
-
+            width = gr.Slider(label="Largura", minimum=512, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+            height = gr.Slider(label="Altura", minimum=512, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+        with gr.Row():
+            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.5, value=7.0)
+            num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=25)
+    gr.Examples(examples=examples, inputs=prompt, outputs=[result, seed], fn=generate, cache_examples=CACHE_EXAMPLES)
+    use_negative_prompt.change(fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt, outputs=negative_prompt, api_name=False)
+    generate_inputs = [prompt, negative_prompt, use_negative_prompt, seed, width, height, guidance_scale, num_inference_steps, randomize_seed]
+    gr.on(triggers=[prompt.submit, run_button.click], fn=generate, inputs=generate_inputs, outputs=[result, seed], api_name="generate_image")
+
+# Lança a interface no ambiente Space
 if __name__ == "__main__":
-    demo.launch()
+    if pipe is not None:
+         print("Lançando interface Gradio no Space...")
+         demo.queue().launch() # Importante para Spaces
+    else:
+         print("ERRO CRÍTICO: Pipeline não carregado ou UNet não substituído corretamente. Lançando UI de erro.")
+         with gr.Blocks() as error_demo:
+              gr.Markdown(f"# Erro ao Carregar Modelo\n{DESCRIPTION}")
+         error_demo.queue().launch()