adding image and mimic

app.py

@@ -11,9 +11,29 @@ from fastapi.openapi.docs import get_swagger_ui_html
 import os
 import requests
 from modules.audio import convert, get_audio_duration
-from modules.r2 import upload_to_s3
+from modules.r2 import upload_to_s3, upload_image_to_s3
 import threading
 import queue
+from diffusers import DiffusionPipeline
+import torch
+from datetime import datetime
+import random
+import numpy as np
+
+SAVE_DIR = "saved_images"  
+if not os.path.exists(SAVE_DIR):
+    os.makedirs(SAVE_DIR, exist_ok=True)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+repo_id = "black-forest-labs/FLUX.1-dev"
+adapter_id = "guardiancc/lora"
+
+pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+pipeline.load_lora_weights(adapter_id)
+pipeline = pipeline.to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
 
 vpv_webhook = os.environ.get("VPV_WEBHOOK")
 
@@ -26,6 +46,34 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+
+def save_generated_image(image):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    unique_id = str(uuid.uuid4())[:8]
+    filename = f"{timestamp}_{unique_id}.png"
+    filepath = os.path.join(SAVE_DIR, filename)
+    image.save(filepath)
+    
+    return filepath
+
+def inference_image(prompt):
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    image = pipeline(
+        prompt=prompt,
+        guidance_scale=3.5,
+        num_inference_steps=20,
+        width=512,
+        height=512,
+        generator=generator,
+        joint_attention_kwargs={"scale": 0.8},
+    ).images[0]
+    
+    filepath = save_generated_image(image, prompt)
+    url = upload_image_to_s3(filepath, os.path.basename(filepath), "png")
+    os.unlink(filepath)
+    return url
+    
 def download_file(url: str) -> str:
     """
     Baixa um arquivo da URL fornecida e o salva no diretório 'downloads/'.
@@ -71,8 +119,67 @@ class ProcessRequest(BaseModel):
     offset: float = -0.3 
     format: str = "wav"
     speed: float = 0.8
-    crossfade: float = 0.06
+    crossfade: float = 0.1
+    
+class ProcessImage(BaseModel):
+    prompt: str
+    id: str
+    receiver: str
+    webhook: str
+
+q = queue.Queue()
+image_queue = queue.Queue()
 
+def process_queue(q):
+    while True:
+        try:
+            key, censor, offset, text, format, speed, crossfade, id, receiver, webhook = q.get(timeout=5)  
+            audio = generate_audio(key, text, censor, offset, speed=speed, crossfade=crossfade) 
+            convertedAudioPath = convert(audio, format)
+            duration = get_audio_duration(convertedAudioPath)
+            audioUrl = upload_to_s3(convertedAudioPath, f"{id}", format)
+            os.remove(audio)
+            os.remove(convertedAudioPath)
+                
+            payload = {
+                    "id": id,
+                    "duration": duration,
+                    "receiver": receiver,
+                    "url": audioUrl
+                }
+                
+            requests.post(webhook, json=payload)
+        except Exception as e:
+            print(e)
+        finally:
+            q.task_done()
+            
+def process_image(q):
+    while True:
+        try:
+            prompt, id, receiver, webhook = q.get(timeout=5)  
+            image = inference_image(prompt)
+                
+            payload = {
+                    "id": id,
+                    "receiver": receiver,
+                    "url": image,
+                    "type": "image"
+                }
+                
+            requests.post(webhook, json=payload)
+        except Exception as e:
+            print(e)
+        finally:
+            q.task_done()
+
+        
+worker_thread = threading.Thread(target=process_queue, args=(q,))
+worker_thread.start()
+
+imagge_worker = threading.Thread(target=process_queue, args=(q,))
+imagge_worker.start()
+        
 @app.post("/process")
 def process_audio(payload: ProcessRequest):
     key = payload.key
@@ -90,21 +197,7 @@ def process_audio(payload: ProcessRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
     try:
-        audio = generate_audio(key, text, censor, offset, speed=speed, crossfade=crossfade) 
-        convertedAudioPath = convert(audio, format)
-        duration = get_audio_duration(convertedAudioPath)
-        audioUrl = upload_to_s3(convertedAudioPath, f"{id}", format)
-        os.remove(audio)
-        os.remove(convertedAudioPath)
-            
-        payload = {
-                "id": id,
-                "duration": duration,
-                "receiver": receiver,
-                "url": audioUrl
-            }
-            
-        requests.post(webhook, json=payload)
+        q.put((key, censor, offset, text, format, speed, crossfade, id, receiver, webhook))
         return {"success": True, "err": ""}
         
     except ValueError as e:
@@ -133,6 +226,47 @@ def process_audio(payload: ProcessRequest):
         requests.post(dc_callback, headers=headers, data=json.dumps(data))
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/image")
+def process_image(payload: ProcessImage):
+    prompt = payload.prompt
+    id = payload.id
+    receiver = payload.receiver
+    webhook = payload.webhook
+    
+    if len(prompt) <= 5:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    try:
+        image_queue.put(( prompt, id, receiver, webhook))
+        return {"success": True, "err": ""}
+        
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    
+    except Exception as e:
+        error_trace = traceback.format_exc() 
+        dc_callback = "https://discord.com/api/webhooks/1285586984898662511/QNVvY2rtoKICamlXsC1BreBaYjS9341jz9ANCDBzayXt4C7v-vTFzKfUtKQkwW7BwpfP"
+        
+        data = {
+            "content": "",
+            "tts": False,
+            "embeds": [
+                {
+                    "type": "rich",
+                    "title": f"Erro aconteceu na IA - MIMIC - 2 ia",
+                    "description": f"Erro: {str(e)}\n\nDetalhes do erro:\n```{error_trace}```"
+                }
+            ]
+        }
+        
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+        requests.post(dc_callback, headers=headers, data=json.dumps(data))
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 class TrainRequest(BaseModel):
     audio: HttpUrl 
     key: str

modules/r2.py

@@ -28,6 +28,15 @@ def upload_to_s3(path, name, extension):
             ExpiresIn=604800
         )
     return url
+
+def upload_image_to_s3(path, name, extension):
+    s3.upload_file(path, bucket, name, ExtraArgs={'ContentType': f'image/{extension}', 'ACL': 'public-read'})
+    url = s3.generate_presigned_url(
+            'get_object',
+            Params={'Bucket': bucket, 'Key': name},
+            ExpiresIn=604800
+        )
+    return url
     
 def get_url(name):
     url = s3.generate_presigned_url(

requirements.txt

@@ -7,4 +7,5 @@ phonemizer
 pydub
 fastapi
 uvicorn
-uuid
+uuid
+diffusers