Update app.py

app.py

@@ -1,15 +1,15 @@
-import spaces  # Import this first to avoid CUDA initialization issues
-import os
 import gradio as gr
 import numpy as np
-import gradio as gr
+import spaces 
 import torch
 import random
 import time
 from PIL import Image
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from huggingface_hub import hf_hub_download
 from gradio_client import Client, handle_file
+import os
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -21,22 +21,45 @@ hf_token = os.getenv('waffles')
 if not hf_token:
     raise ValueError("Hugging Face API token not found. Please set the 'waffles' environment variable.")
 
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
 
-pipe = DiffusionPipeline.from_pretrained("AlekseyCalvin/RCA_Agitprop_Manufactory", torch_dtype=dtype, token=hf_token).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=hf_token).to(device)
+
+@spaces.GPU(duration=300)
+def infer(prompt, seed=0, randomize_seed=True, width=640, height=1024, guidance_scale=0.0, num_inference_steps=5, lora_model="AlekseyCalvin/RCA_Agitprop_Manufactory", progress=gr.Progress(track_tqdm=True)):
+    global pipe
+    
+    # Load LoRA if specified
+    if lora_model:
+        try:
+            pipe.load_lora_weights(lora_model)
+        except Exception as e:
+            return None, seed, f"Failed to load LoRA model: {str(e)}"
 
-@spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-            prompt = prompt, 
-            width = width,
-            height = height,
-            num_inference_steps = num_inference_steps, 
-            generator = generator,
-            guidance_scale=0.0
-    ).images[0] 
+    
+    try:
+        image = pipe(
+            prompt=prompt, 
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps, 
+            generator=generator,
+            guidance_scale=guidance_scale
+        ).images[0]
+        
+        # Unload LoRA weights after generation
+        if lora_model:
+            pipe.unload_lora_weights()
+        
+        return image, seed, "Image generated successfully."
+    except Exception as e:
+        return None, seed, f"Error during image generation: {str(e)}"
+
+    
     return image, seed